Tools/scripts/pickle2db.py

   1 #!/usr/bin/env python
   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 anydbm
  11     -b - open as bsddb btree file
  12     -d - open as dbm file
  13     -g - open as gdbm file
  14     -h - open as bsddb hash file
  15     -r - open as bsddb recno file
  16
  17 The default is hash.  If a pickle file is named it is opened for read
  18 access.  If no pickle file is named, the pickle input is read from standard
  19 input.
  20
  21 Note that recno databases can only contain integer keys, so you can't dump a
  22 hash or btree database using db2pickle.py and reconstitute it to a recno
  23 database with %(prog)s unless your keys are integers.
  24
  25 """
  26
  27 import getopt
  28 try:
  29     import bsddb
  30 except ImportError:
  31     bsddb = None
  32 try:
  33     import dbm
  34 except ImportError:
  35     dbm = None
  36 try:
  37     import gdbm
  38 except ImportError:
  39     gdbm = None
  40 try:
  41     import anydbm
  42 except ImportError:
  43     anydbm = None
  44 import sys
  45 try:
  46     import cPickle as pickle
  47 except ImportError:
  48     import pickle
  49
  50 prog = sys.argv[0]
  51
  52 def usage():
  53     sys.stderr.write(__doc__ % globals())
  54
  55 def main(args):
  56     try:
  57         opts, args = getopt.getopt(args, "hbrdag",
  58                                    ["hash", "btree", "recno", "dbm", "anydbm",
  59                                     "gdbm"])
  60     except getopt.error:
  61         usage()
  62         return 1
  63
  64     if len(args) == 0 or len(args) > 2:
  65         usage()
  66         return 1
  67     elif len(args) == 1:
  68         pfile = sys.stdin
  69         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     dbopen = None
  79     for opt, arg in opts:
  80         if opt in ("-h", "--hash"):
  81             try:
  82                 dbopen = bsddb.hashopen
  83             except AttributeError:
  84                 sys.stderr.write("bsddb module unavailable.\n")
  85                 return 1
  86         elif opt in ("-b", "--btree"):
  87             try:
  88                 dbopen = bsddb.btopen
  89             except AttributeError:
  90                 sys.stderr.write("bsddb module unavailable.\n")
  91                 return 1
  92         elif opt in ("-r", "--recno"):
  93             try:
  94                 dbopen = bsddb.rnopen
  95             except AttributeError:
  96                 sys.stderr.write("bsddb module unavailable.\n")
  97                 return 1
  98         elif opt in ("-a", "--anydbm"):
  99             try:
 100                 dbopen = anydbm.open
 101             except AttributeError:
 102                 sys.stderr.write("anydbm module unavailable.\n")
 103                 return 1
 104         elif opt in ("-g", "--gdbm"):
 105             try:
 106                 dbopen = gdbm.open
 107             except AttributeError:
 108                 sys.stderr.write("gdbm module unavailable.\n")
 109                 return 1
 110         elif opt in ("-d", "--dbm"):
 111             try:
 112                 dbopen = dbm.open
 113             except AttributeError:
 114                 sys.stderr.write("dbm module unavailable.\n")
 115                 return 1
 116     if dbopen is None:
 117         if bsddb is None:
 118             sys.stderr.write("bsddb module unavailable - ")
 119             sys.stderr.write("must specify dbtype.\n")
 120             return 1
 121         else:
 122             dbopen = bsddb.hashopen
 123
 124     try:
 125         db = dbopen(dbfile, 'c')
 126     except bsddb.error:
 127         sys.stderr.write("Unable to open %s.  " % dbfile)
 128         sys.stderr.write("Check for format or version mismatch.\n")
 129         return 1
 130     else:
 131         for k in db.keys():
 132             del db[k]
 133
 134     while 1:
 135         try:
 136             (key, val) = pickle.load(pfile)
 137         except EOFError:
 138             break
 139         db[key] = val
 140
 141     db.close()
 142     pfile.close()
 143
 144     return 0
 145
 146 if __name__ == "__main__":
 147     sys.exit(main(sys.argv[1:]))