Source

fsdict / fsdict.py

Full commit
# −*− coding: UTF−8 −*−
from path import path
import os
import pickle
"""
A class providing dictionary access to a folder.
"""

def get_tmp_dir():
    import tempfile
    return tempfile.mkdtemp()
    
class FSDict(dict):
    """
    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, 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 __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 __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)