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