src/util/cached.py

   1 # This Python file uses the following encoding: utf-8
   2 '''
   3 Created on May 11, 2011
   4
   5 @author: mjacob
   6 '''
   7 import os
   8 import functools
   9 import cPickle
  10
  11 class Cacheable(object):
  12     """interface for a class to have results which are cacheable on disk."""
  13
  14     def __init__(self, cachedir, ignore_cache=None, debug=False):
  15         """
  16         @param cachedir: where to store the cached files on disk. if it does not exist, it will be created.
  17         @param ignore_cache: a collection of components which should not be cached by this instance.
  18         """
  19         self.__cachedir = cachedir
  20         self.__ignore_cache = ignore_cache
  21         self.__debug = debug
  22
  23     def _filename(self, component):
  24         """helper method to generate the filename for a component"""
  25         return os.path.join(self.__cachedir, component)
  26
  27     def is_cached(self, component):
  28         """
  29         returns C{True} if the specified C{component} is cached
  30
  31         @type component: C{str}
  32         @rtype: C{bool}
  33         """
  34         if self.__ignore_cache and component in self.__ignore_cache:
  35             return False
  36
  37         filename = self._filename(component)
  38         found = os.path.exists(filename)
  39
  40         if self.__debug:
  41             if found:
  42                 print "cached file found: %s" % (filename)
  43             else:
  44                 print "cached file not found: %s" % (filename)
  45
  46         return found
  47
  48     def get_cache(self, component, loader=cPickle):
  49         """
  50         returns the specified cached component
  51
  52         @type component: C{str}
  53         @rtype: object
  54         """
  55         filename = self._filename(component)
  56
  57         if self.__debug:
  58             print "reading '%s' from cache" % (filename)
  59
  60         with open(filename) as file:
  61             return loader.load(file)
  62
  63         if self.__debug:
  64             print "finished reading '%s' from cache" % (filename)
  65
  66
  67     def set_cache(self, component, object, dumper=cPickle):
  68         """
  69         writes the specified component to cache.
  70
  71         @type component: C{str}
  72         @type object: anything pickle-able
  73         """
  74
  75         if not os.path.exists(self.__cachedir):
  76             os.makedirs(self.__cachedir)
  77
  78         filename = self._filename(component)
  79
  80         if self.__debug:
  81             print "writing '%s' to cache" % (filename)
  82
  83         with open(filename, 'w') as file:
  84             dumper.dump(object, file)
  85
  86         if self.__debug:
  87             print "finished writing '%s' to cache" % (filename)
  88
  89         return filename
  90
  91
  92 class Cached(object):
  93     """
  94     annotation to indicate a cached function.
  95
  96     using it requires specifying a "key" function with which to generate a filename
  97     to use as the cache for the result of a specific function call.
  98     """
  99
 100     def __init__(self, key=lambda *x, **y: str(x) + str(y), loader=cPickle, dumper=cPickle):
 101         self.__key = key
 102         self.__loader = loader
 103         self.__dumper = dumper
 104
 105     def __call__(self, func):
 106         def caller(*args, **kwargs):
 107             cacher = args[0]
 108             if kwargs:
 109                 component = self.__key(*args[1:], **kwargs)
 110             else:
 111                 component = self.__key(*args[1:])
 112             if cacher.is_cached(component):
 113                 return cacher.get_cache(component, self.__loader)
 114             else:
 115                 if kwargs:
 116                     result = func(*args, **kwargs)
 117                 else:
 118                     result = func(*args)
 119                 cacher.set_cache(component, result, self.__dumper)
 120                 return result
 121         return caller
 122
 123     def __repr__(self):
 124         """Return the function's docstring."""
 125         return self.func.__doc__
 126
 127     def __get__(self, obj, objtype):
 128         """Support instance methods."""
 129         return functools.partial(self.__call__, obj)
 130
 131
 132 class Test(Cacheable):
 133     count = 0
 134
 135     def __init__(self):
 136         Cacheable.__init__(self, os.path.join("/Users/mjacob/testmemoize", "test memoize"))
 137
 138     @Cached(lambda number, *x: str(number))
 139     def do_thing(self, number, *others, **kwargs):
 140         self.count += 1
 141         return (number, self.count)
 142
 143 if __name__ == "__main__":
 144     a = Test()
 145     print a.do_thing(1,4)
 146     print a.do_thing(2,5)
 147     print a.do_thing(3,6)
 148     print a.do_thing(1,4)
 149     print "..."