Commits

dan mackinlay committed 04d88c9

Breaks backward compatibility. Based code of Ian Bicking's example from http://blog.ianbicking.org/2007/08/17/dictmixin/ but subclass to allow pickling of arbitrary values

Comments (0)

Files changed (3)

 # −*− coding: UTF−8 −*−
-from path import path
 import os
 import pickle
+from UserDict import DictMixin
+from shutil import rmtree
+
 """
 A class providing dictionary access to a folder.
+
+My own bloated implementation of thsi has been replace by Ian Bicking's much
+more sleek one, documented with characteristic lucidity at
+http://blog.ianbicking.org/2007/08/17/dictmixin/
 """
 
 def get_tmp_dir():
     import tempfile
     return tempfile.mkdtemp()
 
-    
-class FSDict(dict):
+class FSDict(DictMixin):
     """
     provide dictionary access to a temp dir. Like shelve, but there is no DB
     backing.
     N.B. the keys ordering here is FS-dependent and thus unlike to be the same
     as with a real dict. beware.
     """
+    def __init__(self, path=None, mktmp=False):
+        if path is not None:
+            self.path = path
+        elif mktmp:
+            self.path = get_tmp_dir()
+        else:
+            raise ValueError("no path supplied and no autmatic temp creation")
     
-    def __init__(self, initval=[], work_dir=None, *args, **kwargs):
-        if work_dir is None:
-            work_dir = unicode(get_tmp_dir())
-        self.work_dir = path(work_dir)
-        if not self.work_dir.exists():
-            self.work_dir.mkdir()
-        for key, val in getattr(initval, 'iteritems', initval.__iter__)():
-            self[key] = val
-        super(FSDict, self).__init__(*args, **kwargs)
+    def __getitem__(self, item):
+        fn = os.path.join(self.path, item)
+        if not os.path.exists(fn):
+            raise KeyError("File %s does not exist" % fn)
+        if os.path.isdir(fn):
+            return self.__class__(fn)
+        f = open(fn, 'rb')
+        c = f.read()
+        f.close()
+        return c
     
-    def __contains__(self, key):
-        return (self.work_dir/key).exists()
-        
-    def __setitem__(self, key, val, *args, **kwargs):
-        pickle.dump(val, open(self.work_dir/key, 'wb'))
+    def __setitem__(self, item, value):
+        if item in self:
+            del self[item]
+        fn = os.path.join(self.path, item)
+        if isinstance(value, str):
+            f = open(fn, 'wb')
+            f.write(value)
+            f.close()
+        else:
+            # Assume it is a dictionary
+            os.mkdir(fn)
+            f = self[item]
+            f.update(value)
+
+    def __delitem__(self, item):
+        fn = os.path.join(self.path, item)
+        if not os.path.exists(fn):
+            raise KeyError("File %s does not exist" % fn)
+        if os.path.isdir(fn):
+            ## one way...
+            self[item].clear()
+            os.rmdir(fn)
+            ## another way...
+            #shutil.rmtree(fn)
+        else:
+            os.unlink(fn)
+
+    def keys(self):
+        return os.listdir(self.path)
     
-    def __getitem__(self, key, *args, **kwargs):
-        return pickle.load(open(self.work_dir/key, 'rb'))
-    
-    def __repr__(self):
-        """
-        a hardline list of everything in the dict. may be long.
-        """
-        return repr(dict([(k, v) for k, v in self.iteritems()]))
-        
-    def __str__(self):
-        """
-        str is truncated somewhat.
-        """
-        if len(self.keys()):
-            return '{' + repr(self.keys()[0]) + ':' + repr(self[self.keys()[0]]) + ', ...'
-        else:
-            return super(FSDict, self).__str__()
-    
-    def keys(self, *args, **kwargs):
-        return [key for key in self.iterkeys()]
-    
-    def iterkeys(self, *args, **kwargs):
-        for f in self.work_dir.files():
-            yield str(self.work_dir.relpathto(f))
-        
-    def iteritems(self):
-        for key in self.iterkeys():
-            yield key, self[key]
-            
-    def itervalues(self):
-        for key in self.iterkeys():
-            yield self[key]
-            
-    def __delitem__(self, key, *args, **kwargs):
-        (self.work_dir/key).unlink()
-    
-    def values(self, *args, **kwargs):
-        return [self[key] for key in self.keys()]
-        
-    def cleanup(self):
-        self.work_dir.rmtree()
-    
-    def move(self, new_dir):
-        
-        try:
-            self.work_dir.move(new_dir)
-        except Exception, e:
-            raise
-        else:
-            self.work_dir = new_dir
-    
-    def __eq__(self, other):
-        """
-        when compared to a dict, equate equal if all keys and vals are equal
-        note, this is potentially expensive.
-        """
-        #duck type our way to sanity:
-        if not hasattr(other, 'keys'): return False
-        #OK, it's a dict-ish thing
-        try:
-            return all([self[key]==other[key] for key in other]) and \
-              len(self.keys())==len(other.keys())
-        except KeyError:
-            return False
-    
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        # path objects seem to occasionally resurrect gracelessly
-        state['work_dir'] = unicode(state['work_dir'])
-        return state
-    
-    def __setstate__(self, state):
-        self.__dict__ = state
-        self.work_dir = path(self.work_dir)
-    
-    def setdefault(self, k, d=None):
-        if k not in self:
-            self[k] = d
-        return self[k]
+    def delete(self):
+        rmtree(self.path)
 
+class FSPickleDict(FSDict):
+    """
+    provide dictionary access to a temp dir, storing any picklable object.
+    """
+    def __setitem__(self, item, value):
+        if item in self:
+            del self[item]
+        fn = os.path.join(self.path, item)
+        if isinstance(value, FSDict):
+            os.mkdir(fn)
+            f = self[item]
+            f.update(value)
+        else: #pickle it!
+            with open(fn, 'wb') as f:
+                pickle.dump(value, f)
+                
+    def __getitem__(self, item):
+        fn = os.path.join(self.path, item)
+        if not os.path.exists(fn):
+            raise KeyError("File %s does not exist" % fn)
+        if os.path.isdir(fn):
+            return self.__class__(fn)
+        with open(fn, 'rb') as f:
+            return pickle.load(f)
+
 from setuptools import setup, find_packages
 import sys, os
 
-version = '0.1.1'
+version = '0.2alpha'
 
 setup(name='fsdict',
       version=version,
       include_package_data=True,
       zip_safe=True,
       install_requires=[
-          'path.py' 
       ],
       entry_points="""
       # -*- Entry points: -*-

tests/fsdict_tests.py

 from __future__ import with_statement
 from nose import with_setup
 from nose.tools import nottest, assert_equal, assert_true, assert_false
-from fsdict import FSDict
+from fsdict import FSPickleDict
 import pickle
 from pickle import loads, dumps
+import os
+import os.path
+from os.path import join as pathjoin
+from os.path import exists as pathexists
 
 def test_fsdict_pickles():
-    "occasionally something goes wrong with work_dir in pickling"
-    d1 = FSDict({'a':'b'})
+    "occasionally something goes wrong with path in pickling"
+    d1 = FSPickleDict(mktmp=True)
+    d1.update({'a':'b'})
     d2 = loads(dumps(d1))
     yield assert_equal, d1, d2
-    yield assert_equal, d1.work_dir, d2.work_dir
+    yield assert_equal, d1.path, d2.path
 
 def test_fsdict_loads():
     import pickle
     
-    d = FSDict({'a':'b'})
-    yield assert_true, (d.work_dir).exists()
-    yield assert_true, (d.work_dir/'a').exists()
+    d = FSPickleDict(mktmp=True)
+    d.update({'a':'b'})
+    yield assert_true, pathexists(d.path)
+    yield assert_true, pathexists(pathjoin(d.path, 'a'))
     yield assert_equal, d['a'], 'b'
     yield assert_true, 'a' in d
     yield assert_false, 'q' in d
         yield assert_equal, val, 'b'
     yield assert_equal, d.values(), ['b']
     yield assert_equal, d['a'], 'b'
-    pickle.dump('c', open(d.work_dir/'a', 'w'))
+    subdir = pathjoin(d.path, 'a')
+    # os.makedirs(pathjoin(d.path, 'd'))
+    # os.makedirs(subdir)
+    pickle.dump('c', open(subdir, 'w'))
     yield assert_equal, d['a'], 'c'
     del(d['a'])
-    yield assert_false, (d.work_dir/'a').exists()
-    d.cleanup()
-    yield assert_false, (d.work_dir).exists()
-    e=FSDict()
-    yield assert_true, e.work_dir.exists()
-    f=FSDict()
-    yield assert_true, f.work_dir.exists()
-    e.cleanup()
-    f.cleanup()
-    yield assert_false, e.work_dir.exists()
-    yield assert_false, f.work_dir.exists()
+    yield assert_false, pathexists(pathjoin(d.path, 'a'))
+    d.delete()
+    yield assert_false, pathexists(d.path)
+    e=FSPickleDict(mktmp=True)
+    yield assert_true, pathexists(e.path)
+    f=FSPickleDict(mktmp=True)
+    yield assert_true, pathexists(f.path)
+    e.delete()
+    f.delete()
+    yield assert_false, pathexists(e.path)
+    yield assert_false, pathexists(f.path)
 
 def test_fsdict_equals():
     """
     as far as equality goes, we want these things to equate as normal dicts
     """
-    fsd = FSDict({'a':'b'})
+    fsd = FSPickleDict(mktmp=True)
+    fsd.update({'a':'b'})
     d1 = {'a':'b'}
     d2 = {'c':'b'}
     d3 = {'a':'c'}
     
 def test_fsdict_setdefault():
     """does setdefault work as expected?"""
-    d = FSDict()
+    d = FSPickleDict(mktmp=True)
     yield assert_equal, d.setdefault('a',2), 2
     #sure, but did that REALLY go to the FS?
-    yield assert_true, (d.work_dir/'a').exists()
+    yield assert_true, pathexists(pathjoin(d.path, 'a'))
     yield assert_equal, d.setdefault('a',3), 2
-    d.cleanup()
+    d.delete()