pickle2db.py 3.93 KB
Newer Older
1
#!/usr/bin/env python3
2 3

"""
4
Synopsis: %(prog)s [-h|-b|-g|-r|-a|-d] [ picklefile ] dbfile
5 6

Read the given picklefile as a series of key/value pairs and write to a new
7
database.  If the database already exists, any contents are deleted.  The
8
optional flags indicate the type of the output database:
9

10
    -a - open using dbm (open any supported format)
11
    -b - open as bsddb btree file
12 13
    -d - open as dbm.ndbm file
    -g - open as dbm.gnu file
14 15 16 17 18 19 20 21
    -h - open as bsddb hash file
    -r - open as bsddb recno file

The default is hash.  If a pickle file is named it is opened for read
access.  If no pickle file is named, the pickle input is read from standard
input.

Note that recno databases can only contain integer keys, so you can't dump a
Skip Montanaro's avatar
Skip Montanaro committed
22
hash or btree database using db2pickle.py and reconstitute it to a recno
23 24
database with %(prog)s unless your keys are integers.

25 26 27 28 29 30 31 32
"""

import getopt
try:
    import bsddb
except ImportError:
    bsddb = None
try:
33
    import dbm.ndbm as dbm
34 35
except ImportError:
    dbm = None
36
try:
37
    import dbm.gnu as gdbm
38 39
except ImportError:
    gdbm = None
40
try:
Benjamin Peterson's avatar
Benjamin Peterson committed
41
    import dbm.ndbm as anydbm
42 43 44 45
except ImportError:
    anydbm = None
import sys
try:
Benjamin Peterson's avatar
Benjamin Peterson committed
46
    import pickle as pickle
47 48 49 50 51 52
except ImportError:
    import pickle

prog = sys.argv[0]

def usage():
Tim Peters's avatar
Tim Peters committed
53
    sys.stderr.write(__doc__ % globals())
54 55 56

def main(args):
    try:
57 58 59
        opts, args = getopt.getopt(args, "hbrdag",
                                   ["hash", "btree", "recno", "dbm", "anydbm",
                                    "gdbm"])
60 61 62 63 64 65 66 67 68 69
    except getopt.error:
        usage()
        return 1

    if len(args) == 0 or len(args) > 2:
        usage()
        return 1
    elif len(args) == 1:
        pfile = sys.stdin
        dbfile = args[0]
70
    else:
71
        try:
72
            pfile = open(args[0], 'rb')
73
        except IOError:
74
            sys.stderr.write("Unable to open %s\n" % args[0])
75
            return 1
76
        dbfile = args[1]
77 78 79 80 81 82 83

    dbopen = None
    for opt, arg in opts:
        if opt in ("-h", "--hash"):
            try:
                dbopen = bsddb.hashopen
            except AttributeError:
84
                sys.stderr.write("bsddb module unavailable.\n")
85 86 87 88 89
                return 1
        elif opt in ("-b", "--btree"):
            try:
                dbopen = bsddb.btopen
            except AttributeError:
90
                sys.stderr.write("bsddb module unavailable.\n")
91 92 93 94 95
                return 1
        elif opt in ("-r", "--recno"):
            try:
                dbopen = bsddb.rnopen
            except AttributeError:
96
                sys.stderr.write("bsddb module unavailable.\n")
97 98 99 100 101
                return 1
        elif opt in ("-a", "--anydbm"):
            try:
                dbopen = anydbm.open
            except AttributeError:
102
                sys.stderr.write("dbm module unavailable.\n")
103
                return 1
104 105 106 107
        elif opt in ("-g", "--gdbm"):
            try:
                dbopen = gdbm.open
            except AttributeError:
108
                sys.stderr.write("dbm.gnu module unavailable.\n")
109
                return 1
110 111 112 113
        elif opt in ("-d", "--dbm"):
            try:
                dbopen = dbm.open
            except AttributeError:
114
                sys.stderr.write("dbm.ndbm module unavailable.\n")
115 116 117
                return 1
    if dbopen is None:
        if bsddb is None:
118 119
            sys.stderr.write("bsddb module unavailable - ")
            sys.stderr.write("must specify dbtype.\n")
120 121 122 123 124 125 126
            return 1
        else:
            dbopen = bsddb.hashopen

    try:
        db = dbopen(dbfile, 'c')
    except bsddb.error:
127 128
        sys.stderr.write("Unable to open %s.  " % dbfile)
        sys.stderr.write("Check for format or version mismatch.\n")
129 130
        return 1
    else:
131
        for k in list(db.keys()):
132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147
            del db[k]

    while 1:
        try:
            (key, val) = pickle.load(pfile)
        except EOFError:
            break
        db[key] = val

    db.close()
    pfile.close()

    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))