pickle2db.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. #!/usr/bin/env python3
  2. """
  3. Synopsis: %(prog)s [-h|-b|-g|-r|-a|-d] [ picklefile ] dbfile
  4. Read the given picklefile as a series of key/value pairs and write to a new
  5. database. If the database already exists, any contents are deleted. The
  6. optional flags indicate the type of the output database:
  7. -a - open using dbm (open any supported format)
  8. -b - open as bsddb btree file
  9. -d - open as dbm.ndbm file
  10. -g - open as dbm.gnu file
  11. -h - open as bsddb hash file
  12. -r - open as bsddb recno file
  13. The default is hash. If a pickle file is named it is opened for read
  14. access. If no pickle file is named, the pickle input is read from standard
  15. input.
  16. Note that recno databases can only contain integer keys, so you can't dump a
  17. hash or btree database using db2pickle.py and reconstitute it to a recno
  18. database with %(prog)s unless your keys are integers.
  19. """
  20. import getopt
  21. try:
  22. import bsddb
  23. except ImportError:
  24. bsddb = None
  25. try:
  26. import dbm.ndbm as dbm
  27. except ImportError:
  28. dbm = None
  29. try:
  30. import dbm.gnu as gdbm
  31. except ImportError:
  32. gdbm = None
  33. try:
  34. import dbm.ndbm as anydbm
  35. except ImportError:
  36. anydbm = None
  37. import sys
  38. try:
  39. import pickle as pickle
  40. except ImportError:
  41. import pickle
  42. prog = sys.argv[0]
  43. def usage():
  44. sys.stderr.write(__doc__ % globals())
  45. def main(args):
  46. try:
  47. opts, args = getopt.getopt(args, "hbrdag",
  48. ["hash", "btree", "recno", "dbm", "anydbm",
  49. "gdbm"])
  50. except getopt.error:
  51. usage()
  52. return 1
  53. if len(args) == 0 or len(args) > 2:
  54. usage()
  55. return 1
  56. elif len(args) == 1:
  57. pfile = sys.stdin
  58. dbfile = args[0]
  59. else:
  60. try:
  61. pfile = open(args[0], 'rb')
  62. except IOError:
  63. sys.stderr.write("Unable to open %s\n" % args[0])
  64. return 1
  65. dbfile = args[1]
  66. dbopen = None
  67. for opt, arg in opts:
  68. if opt in ("-h", "--hash"):
  69. try:
  70. dbopen = bsddb.hashopen
  71. except AttributeError:
  72. sys.stderr.write("bsddb module unavailable.\n")
  73. return 1
  74. elif opt in ("-b", "--btree"):
  75. try:
  76. dbopen = bsddb.btopen
  77. except AttributeError:
  78. sys.stderr.write("bsddb module unavailable.\n")
  79. return 1
  80. elif opt in ("-r", "--recno"):
  81. try:
  82. dbopen = bsddb.rnopen
  83. except AttributeError:
  84. sys.stderr.write("bsddb module unavailable.\n")
  85. return 1
  86. elif opt in ("-a", "--anydbm"):
  87. try:
  88. dbopen = anydbm.open
  89. except AttributeError:
  90. sys.stderr.write("dbm module unavailable.\n")
  91. return 1
  92. elif opt in ("-g", "--gdbm"):
  93. try:
  94. dbopen = gdbm.open
  95. except AttributeError:
  96. sys.stderr.write("dbm.gnu module unavailable.\n")
  97. return 1
  98. elif opt in ("-d", "--dbm"):
  99. try:
  100. dbopen = dbm.open
  101. except AttributeError:
  102. sys.stderr.write("dbm.ndbm module unavailable.\n")
  103. return 1
  104. if dbopen is None:
  105. if bsddb is None:
  106. sys.stderr.write("bsddb module unavailable - ")
  107. sys.stderr.write("must specify dbtype.\n")
  108. return 1
  109. else:
  110. dbopen = bsddb.hashopen
  111. try:
  112. db = dbopen(dbfile, 'c')
  113. except bsddb.error:
  114. sys.stderr.write("Unable to open %s. " % dbfile)
  115. sys.stderr.write("Check for format or version mismatch.\n")
  116. return 1
  117. else:
  118. for k in list(db.keys()):
  119. del db[k]
  120. while 1:
  121. try:
  122. (key, val) = pickle.load(pfile)
  123. except EOFError:
  124. break
  125. db[key] = val
  126. db.close()
  127. pfile.close()
  128. return 0
  129. if __name__ == "__main__":
  130. sys.exit(main(sys.argv[1:]))