[Python-checkins] r69327 - sandbox/trunk/dbm_sqlite/alt/dbdict.py

[raymond.hettinger]

Author: raymond.hettinger
Date: Fri Feb  6 01:38:48 2009
New Revision: 69327

Log:
Auto search input file formats.

Modified:
   sandbox/trunk/dbm_sqlite/alt/dbdict.py

Modified: sandbox/trunk/dbm_sqlite/alt/dbdict.py
==============================================================================
--- sandbox/trunk/dbm_sqlite/alt/dbdict.py	(original)
+++ sandbox/trunk/dbm_sqlite/alt/dbdict.py	Fri Feb  6 01:38:48 2009
@@ -5,7 +5,8 @@
 close time's are potentially long because the whole dict must be
 read or written to disk.
 
-File format is selectable between pickle, json, eval, and csv.
+Input file format is automatically discovered.
+Output file format is selectable between pickle, json, and csv.
 All three are backed by fast C implementations.
 
 '''
@@ -15,16 +16,18 @@
 
 class DictDB(dict):
 
-    def __init__(self, filename, flag='c', mode=None, format='csv', *args, **kwds):
-        # r=readonly   c=read_write_create_if_needed   n=new_overwrite_previous
+    def __init__(self, filename, flag='c', mode=None, format=None, *args, **kwds):
+        # flags:    r=readonly   c=create   n=new_overwrite_previous
+        # mode:     None or a triple like 0x666
+        # format:  csv, json, or pickle (default is csv)
         self.flag = flag
         self.mode = mode
-        self.format = format
+        self.format = format or 'csv'
         self.filename = filename
         if flag != 'n' and os.access(filename, os.R_OK):
             file = __builtins__.open(filename, 'rb')
             try:
-                self.update(self.load(file))
+                self.load(file)
             finally:
                 file.close()
         self.update(*args, **kwds)
@@ -42,40 +45,34 @@
                 raise
             file.close()
             shutil.move(file.name, self.filename)    # atomic commit
-            if self.mode:
+            if self.mode is not None:
                 os.chmod(self.filename, self.mode)
 
     def close(self):
         self.sync()
 
-    def dump(self, f):
+    def dump(self, file):
         if self.format == 'csv':
-            csv.writer(f).writerows(self.iteritems())
+            csv.writer(file).writerows(self.iteritems())
         elif self.format == 'json':
-            json.dump(self, f, separators=(',',':'))
+            json.dump(self, file, separators=(',',':'))
         elif self.format == 'pickle':
-            pickle.dump(self.items(), f, -1)
-        elif self.format == 'eval':
-            f.write(repr(self))
+            pickle.dump(self.items(), file, -1)
         else:
             raise NotImplementedError('Unknown format: %r' % self.format)
 
-    def load(self, f):
-        # Use csv or json for untrusted data.  Handles only basic types.
-        # Use pickle or eval only for trusted.  Handles many object types.
-        if self.format == 'csv':
-            return csv.reader(f)
-        elif self.format == 'json':
-            return json.load(f)
-        elif self.format == 'pickle':
-            return pickle.load(f)
-        elif self.format == 'eval':
-            return eval(f.read())
-        else:
-            raise NotImplementedError('Unknown format: %r' % self.format)
+    def load(self, file):
+        # try formats from most restrictive to least restrictive
+        for loader in (pickle.load, json.load, csv.reader):
+            file.seek(0)
+            try:
+                return self.update(loader(file))
+            except Exception:
+                pass
+        raise NotImplementedError('Unknown format: %r' % self.format)
 
 
-def dbopen(filename, flag='c', mode=0o666, format='csv'):
+def dbopen(filename, flag='c', mode=None, format=None):
     return DictDB(filename, flag, mode, format)