1# Copyright 2014 The Chromium Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5""" 6Collection of decorators to help optimize Python functions and classes 7by caching return values. 8 9Return values are (by default) keyed off of the function's parameters. 10Consequently, any parameters that are used in memoization must be hashable. 11 12This library offers two styles of memoization: 131) Absolute memoization (memo) uses the full set of parameters against a 14 per-function memo dictionary. If this is used on an instance method, the 15 'self' parameter is included in the key. 162) Instance memoization (memo_i) uses a per-instance memoization dictionary. 17 The dictionary is stored as a member ('_memo__dict') of of the instance. 18 Consequently, the 'self' parameter is no longer needed/used in the 19 memoization key, removing the need to have the instance itself support being 20 hashed. 21 22Memoized function state can be cleared by calling the memoized function's 23'memo_clear' method. 24""" 25 26import inspect 27 28 29# Instance variable added to class instances that use memoization to 30# differentiate their memoization values. 31MEMO_INSTANCE_VARIABLE = '_memo__dict' 32 33 34class MemoizedFunction(object): 35 """Handles the memoization state of a given memoized function.""" 36 37 # Memoization constant to represent 'un-memoized' (b/c 'None' is a valid 38 # result) 39 class _EMPTY(object): 40 pass 41 EMPTY = _EMPTY() 42 43 def __init__(self, func, ignore=None, memo_dict=None): 44 """ 45 Args: 46 func: (function) The function to memoize 47 ignore: (container) The names of 'func' parameters to ignore when 48 generating its memo key. Only parameters that have no effect on the 49 output of the function should be included. 50 """ 51 self.func = func 52 self.ignore = frozenset(ignore or ()) 53 self.memo_dict = memo_dict 54 self.im_self = None 55 self.im_class = None 56 57 def __repr__(self): 58 properties = [str(self.func)] 59 if self.im_self is not None: 60 properties.append('bound=%s' % (self.im_self,)) 61 if len(self.ignore) > 0: 62 properties.append('ignore=%s' % (','.join(sorted(self.ignore)))) 63 return '%s(%s)' % (type(self).__name__, ', '.join(properties)) 64 65 def __get__(self, obj, klass=None): 66 # Make this callable class a bindable Descriptor 67 if klass is None: 68 klass = type(obj) 69 self.im_self = obj 70 self.im_class = klass 71 return self 72 73 def _get_call_args(self, args): 74 """Returns the call arguments, factoring in 'self' if this method is bound. 75 """ 76 if self.im_self is not None: 77 return (self.im_self,) + args 78 return args 79 80 def _get_memo_dict(self): 81 """Returns: (dict) the memoization dictionary to store return values in.""" 82 memo_dict = None 83 if self.im_self is not None: 84 # Is the instance dictionary defined? 85 memo_dict = getattr(self.im_self, MEMO_INSTANCE_VARIABLE, None) 86 if memo_dict is None: 87 memo_dict = {} 88 setattr(self.im_self, MEMO_INSTANCE_VARIABLE, memo_dict) 89 return memo_dict 90 91 # No instance dict; use our local 'memo_dict'. 92 if self.memo_dict is None: 93 self.memo_dict = {} 94 return self.memo_dict 95 96 def _key(self, args, kwargs): 97 """Returns: the memoization key for a set of function arguments. 98 99 This 'ignored' parameters are removed prior to generating the key. 100 """ 101 if self.im_self is not None: 102 # We are bound to an instance; use None for args[0] ("self"). 103 args = (None,) + tuple(args) 104 105 call_params = inspect.getcallargs(self.func, *args, **kwargs) 106 return tuple((k, v) 107 for k, v in sorted(call_params.iteritems()) 108 if k not in self.ignore) 109 110 def __call__(self, *args, **kwargs): 111 """Retrieves the memoized function result. 112 113 If the memoized function has not been memoized, it will be invoked; 114 otherwise, the memoized value will be returned. 115 116 Args: 117 memo_key: The generated memoization key for this invocation. 118 args, kwargs: Function parameters (only used if not memoized yet) 119 Returns: 120 The memoized function's return value. 121 """ 122 123 memo_dict = self._get_memo_dict() 124 memo_key = self._key(args, kwargs) 125 result = memo_dict.get(memo_key, self.EMPTY) 126 if result is self.EMPTY: 127 args = self._get_call_args(args) 128 result = memo_dict[memo_key] = self.func(*args, **kwargs) 129 return result 130 131 def memo_clear(self, *args, **kwargs): 132 """Clears memoization results for a given set of arguments. 133 134 If no arguments are supplied, this will clear all retained memoization for 135 this function. 136 137 If no memoized result is stored for the supplied parameters, this function 138 is a no-op. 139 140 Args: 141 args, kwargs: Memoization function parameters whose memoized value should 142 be cleared. 143 """ 144 memo_dict = self._get_memo_dict() 145 146 if args or kwargs: 147 memo_key = self._key(args, kwargs) 148 memo_dict.pop(memo_key, None) 149 else: 150 memo_dict.clear() 151 152 153 154def memo(ignore=None, memo_dict=None): 155 """Generic function memoization decorator. 156 157 This memoizes a specific function using a function key. 158 159 The following example memoizes the function absolutely. It will be executed 160 once and, after that, will only returned cached results. 161 162 @memo.memo(ignore=('print_debug_output',)) 163 def my_method(print_debug_output=False): 164 # Perform complex calculation 165 if print_debug_output: 166 print 'The calculated result is: %r' % (result) 167 return result 168 169 The following example memoizes for unique values of 'param1' and 'param2', 170 but not 'print_debug_output' since it doesn't affect the function's result: 171 172 @memo.memo() 173 def my_method(param1, param2): 174 # Perform complex calculation using 'param1' and 'param2' 175 if print_debug_output: 176 print 'The calculated result for (%r, %r) is: %r' % 177 (param1, param2, result) 178 return result 179 180 Args: 181 ignore: (list) The names of parameters to ignore when memoizing. 182 """ 183 def wrap(func): 184 return MemoizedFunction( 185 func, 186 ignore=ignore, 187 memo_dict=memo_dict, 188 ) 189 return wrap 190