Source

fsdict / fsdict.py

Full commit
# −*− coding: UTF−8 −*−
import os
import pickle
from UserDict import DictMixin
from shutil import rmtree

"""
A class providing dictionary access to a folder.

My own original bloated implementation of this has been replaced by Ian
Bicking's much more sleek one, documented with characteristic lucidity at
http://blog.ianbicking.org/2007/08/17/dictmixin/

For my purposes, automatic tmpdir creation and automatic GZipping are helpful, so I've added these features.
"""

def _get_tmp_dir():
    import tempfile
    return tempfile.mkdtemp()

class FSDict(DictMixin):
    """
    provide dictionary access to a temp dir. Like shelve, but there is no DB
    backing, and it can, e.g. rsync.
    
    N.B. the keys ordering here is FS-dependent and thus unlike to be the same
    as with a real dict. Beware.
    
    If you pass "handle_factory=gzip.GzipFile", output will be gzipped. Or you
    can pass any other kind of handle factory to access other exotic kinds of
    FS-like storage.
    """
    def __init__(self, path=None, mktmp=False, handle_factory=open):
        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")
        self._handle_factory = handle_factory
    
    def __getitem__(self, item):
        fn = self._get_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 = self._handle_factory(fn, 'rb')
        c = f.read()
        f.close()
        return c
    
    def __setitem__(self, item, value):
        if item in self:
            del self[item]
        fn = self._get_path(item)
        if isinstance(value, str):
            f = self._handle_factory(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 = self._get_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 delete(self):
        rmtree(self.path)
    
    def _get_path(self, file_name):
        return os.path.join(self.path, file_name)

class FSPickleDict(FSDict):
    """
    Provide dictionary access to a temp dir, storing any picklable object.
    Note that this is not completely symmtrical with FSDict, which will store
    *only* strings in values; this will store *anything but* FSDicts in
    values.
    """
    def __setitem__(self, item, value):
        if item in self:
            del self[item]
        fn = self._get_path(item)
        if isinstance(value, FSDict):
            os.mkdir(fn)
            f = self[item]
            f.update(value)
        else: #pickle it!
            with self._handle_factory(fn, 'wb') as f:
                pickle.dump(value, f)
                
    def __getitem__(self, item):
        fn = self._get_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 self._handle_factory(fn, 'rb') as f:
            return pickle.load(f)