sympy/core/cache.py

   1 """ Caching facility for SymPy """
   2
   3 # TODO: refactor CACHE & friends into class?
   4
   5 # global cache registry:
   6 CACHE = []  # [] of
   7             #    (item, {} or tuple of {})
   8
   9 def print_cache():
  10     """print cache content"""
  11
  12     for item, cache in CACHE:
  13         item = str(item)
  14         head = '='*len(item)
  15
  16         print head
  17         print item
  18         print head
  19
  20         if not isinstance(cache, tuple):
  21             cache = (cache,)
  22             shown = False
  23         else:
  24             shown = True
  25
  26         for i, kv in enumerate(cache):
  27             if shown:
  28                 print '\n*** %i ***\n' % i
  29
  30             for k, v in kv.iteritems():
  31                 print '  %s :\t%s' % (k, v)
  32
  33 def clear_cache():
  34     """clear cache content"""
  35     for item, cache in CACHE:
  36         if not isinstance(cache, tuple):
  37             cache = (cache,)
  38
  39         for kv in cache:
  40             kv.clear()
  41
  42 ########################################
  43
  44 def __cacheit_nocache(func):
  45     return func
  46
  47 def __cacheit(func):
  48     """caching decorator.
  49
  50        important: the result of cached function must be *immutable*
  51
  52
  53        Example
  54        -------
  55
  56        @cacheit
  57        def f(a,b):
  58            return a+b
  59
  60
  61        @cacheit
  62        def f(a,b):
  63            return [a,b] # <-- WRONG, returns mutable object
  64
  65
  66        to force cacheit to check returned results mutability and consistency,
  67        set environment variable SYMPY_USE_CACHE to 'debug'
  68     """
  69
  70     func._cache_it_cache = func_cache_it_cache = {}
  71     CACHE.append((func, func_cache_it_cache))
  72
  73     def wrapper(*args, **kw_args):
  74         if kw_args:
  75             keys = kw_args.keys()
  76             keys.sort()
  77             items = [(k+'=',kw_args[k]) for k in keys]
  78             k = args + tuple(items)
  79         else:
  80             k = args
  81         try:
  82             return func_cache_it_cache[k]
  83         except KeyError:
  84             pass
  85         func_cache_it_cache[k] = r = func(*args, **kw_args)
  86         return r
  87
  88     wrapper.__doc__ = func.__doc__
  89     wrapper.__name__ = func.__name__
  90
  91     return wrapper
  92
  93 def __cacheit_debug(func):
  94     """cacheit + code to check cache consitency"""
  95     cfunc = __cacheit(func)
  96
  97     def wrapper(*args, **kw_args):
  98         # always call function itself and compare it with cached version
  99         r1 = func (*args, **kw_args)
 100         r2 = cfunc(*args, **kw_args)
 101
 102         # try to see if the result is immutable
 103         #
 104         # this works because:
 105         #
 106         # hash([1,2,3])         -> raise TypeError
 107         # hash({'a':1, 'b':2})  -> raise TypeError
 108         # hash((1,[2,3]))       -> raise TypeError
 109         #
 110         # hash((1,2,3))         -> just computes the hash
 111         hash(r1), hash(r2)
 112
 113         # also see if returned values are the same
 114         assert r1 == r2
 115
 116         return r1
 117
 118     wrapper.__doc__ = func.__doc__
 119     wrapper.__name__ = func.__name__
 120
 121     return wrapper
 122
 123 def __cacheit_nondummy(func):
 124     func._cache_it_cache = func_cache_it_cache = {}
 125     CACHE.append((func, func_cache_it_cache))
 126
 127     def wrapper(*args, **kw_args):
 128         if kw_args:
 129             try:
 130                 dummy = kw_args['dummy']
 131             except KeyError:
 132                 dummy = None
 133             if dummy:
 134                 return func(*args, **kw_args)
 135             keys = kw_args.keys()
 136             keys.sort()
 137             items = [(k+'=',kw_args[k]) for k in keys]
 138             k = args + tuple(items)
 139         else:
 140             k = args
 141         try:
 142             return func_cache_it_cache[k]
 143         except KeyError:
 144             pass
 145         func_cache_it_cache[k] = r = func(*args, **kw_args)
 146         return r
 147
 148     wrapper.__doc__ = func.__doc__
 149     wrapper.__name__ = func.__name__
 150
 151     return wrapper
 152
 153 class MemoizerArg:
 154     """ See Memoizer.
 155     """
 156
 157     def __init__(self, allowed_types, converter = None, name = None):
 158         self._allowed_types = allowed_types
 159         self.converter = converter
 160         self.name = name
 161
 162     def fix_allowed_types(self, have_been_here={}):
 163         i = id(self)
 164         if have_been_here.get(i): return
 165         allowed_types = self._allowed_types
 166         if isinstance(allowed_types, str):
 167             self.allowed_types = getattr(C, allowed_types)
 168         elif isinstance(allowed_types, (tuple, list)):
 169             new_allowed_types = []
 170             for t in allowed_types:
 171                 if isinstance(t, str):
 172                     t = getattr(C, t)
 173                 new_allowed_types.append(t)
 174             self.allowed_types = tuple(new_allowed_types)
 175         else:
 176             self.allowed_types = allowed_types
 177         have_been_here[i] = True
 178         return
 179
 180     def process(self, obj, func, index = None):
 181         if isinstance(obj, self.allowed_types):
 182             if self.converter is not None:
 183                 obj = self.converter(obj)
 184             return obj
 185         func_src = '%s:%s:function %s' % (func.func_code.co_filename, func.func_code.co_firstlineno, func.func_name)
 186         if index is None:
 187             raise ValueError('%s return value must be of type %r but got %r' % (func_src, self.allowed_types, obj))
 188         if isinstance(index, (int,long)):
 189             raise ValueError('%s %s-th argument must be of type %r but got %r' % (func_src, index, self.allowed_types, obj))
 190         if isinstance(index, str):
 191             raise ValueError('%s %r keyword argument must be of type %r but got %r' % (func_src, index, self.allowed_types, obj))
 192         raise NotImplementedError(`index,type(index)`)
 193
 194 class Memoizer:
 195     """ Memoizer function decorator generator.
 196
 197     Features:
 198       - checks that function arguments have allowed types
 199       - optionally apply converters to arguments
 200       - cache the results of function calls
 201       - optionally apply converter to function values
 202
 203     Usage:
 204
 205       @Memoizer(<allowed types for argument 0>,
 206                 MemoizerArg(<allowed types for argument 1>),
 207                 MemoizerArg(<allowed types for argument 2>, <convert argument before function call>),
 208                 MemoizerArg(<allowed types for argument 3>, <convert argument before function call>, name=<kw argument name>),
 209                 ...
 210                 return_value_converter = <None or converter function, usually makes a copy>
 211                 )
 212       def function(<arguments>, <kw_argumnets>):
 213           ...
 214
 215     Details:
 216       - if allowed type is string object then there `C` must have attribute
 217         with the string name that is used as the allowed type --- this is needed
 218         for applying Memoizer decorator to Basic methods when Basic definition
 219         is not defined.
 220
 221     Restrictions:
 222       - arguments must be immutable
 223       - when function values are mutable then one must use return_value_converter to
 224         deep copy the returned values
 225
 226     Ref: http://en.wikipedia.org/wiki/Memoization
 227     """
 228
 229     def __init__(self, *arg_templates, **kw_arg_templates):
 230         new_arg_templates = []
 231         for t in arg_templates:
 232             if not isinstance(t, MemoizerArg):
 233                 t = MemoizerArg(t)
 234             new_arg_templates.append(t)
 235         self.arg_templates = tuple(new_arg_templates)
 236         return_value_converter = kw_arg_templates.pop('return_value_converter', None)
 237         self.kw_arg_templates = kw_arg_templates.copy()
 238         for template in self.arg_templates:
 239             if template.name is not None:
 240                 self.kw_arg_templates[template.name] = template
 241         if return_value_converter is None:
 242             self.return_value_converter = lambda obj: obj
 243         else:
 244             self.return_value_converter = return_value_converter
 245
 246     def fix_allowed_types(self, have_been_here={}):
 247         i = id(self)
 248         if have_been_here.get(i): return
 249         for t in self.arg_templates:
 250             t.fix_allowed_types()
 251         for k,t in self.kw_arg_templates.items():
 252             t.fix_allowed_types()
 253         have_been_here[i] = True
 254
 255     def __call__(self, func):
 256         cache = {}
 257         value_cache = {}
 258         CACHE.append((func, (cache, value_cache)))
 259
 260         def wrapper(*args, **kw_args):
 261             kw_items = tuple(kw_args.items())
 262             try:
 263                 return self.return_value_converter(cache[args,kw_items])
 264             except KeyError:
 265                 pass
 266             self.fix_allowed_types()
 267             new_args = tuple([template.process(a,func,i) for (a, template, i) in zip(args, self.arg_templates, range(len(args)))])
 268             assert len(args)==len(new_args)
 269             new_kw_args = {}
 270             for k, v in kw_items:
 271                 template = self.kw_arg_templates[k]
 272                 v = template.process(v, func, k)
 273                 new_kw_args[k] = v
 274             new_kw_items = tuple(new_kw_args.items())
 275             try:
 276                 return self.return_value_converter(cache[new_args, new_kw_items])
 277             except KeyError:
 278                 r = func(*new_args, **new_kw_args)
 279                 try:
 280                     try:
 281                         r = value_cache[r]
 282                     except KeyError:
 283                         value_cache[r] = r
 284                 except TypeError:
 285                     pass
 286                 cache[new_args, new_kw_items] = cache[args, kw_items] = r
 287                 return self.return_value_converter(r)
 288         return wrapper
 289
 290
 291 class Memoizer_nocache(Memoizer):
 292
 293     def __call__(self, func):
 294         # XXX I would be happy just to return func, but we need to provide
 295         # argument convertion, and it is really needed for e.g. Real("0.5")
 296         def wrapper(*args, **kw_args):
 297             kw_items = tuple(kw_args.items())
 298             self.fix_allowed_types()
 299             new_args = tuple([template.process(a,func,i) for (a, template, i) in zip(args, self.arg_templates, range(len(args)))])
 300             assert len(args)==len(new_args)
 301             new_kw_args = {}
 302             for k, v in kw_items:
 303                 template = self.kw_arg_templates[k]
 304                 v = template.process(v, func, k)
 305                 new_kw_args[k] = v
 306
 307             r = func(*new_args, **new_kw_args)
 308             return self.return_value_converter(r)
 309
 310         return wrapper
 311
 312
 313
 314 # SYMPY_USE_CACHE=yes/no/debug
 315 import os
 316 usecache = os.getenv('SYMPY_USE_CACHE', 'yes').lower()
 317
 318 if usecache=='no':
 319     Memoizer            = Memoizer_nocache
 320     cacheit             = __cacheit_nocache
 321 elif usecache=='yes':
 322     cacheit = __cacheit
 323 elif usecache=='debug':
 324     cacheit = __cacheit_debug   # a lot slower
 325 else:
 326     raise RuntimeError('unknown argument in SYMPY_USE_CACHE: %s' % usecache)