blob: ae4bbdfc03c287afecdc4e10fd4b72ac68ce5275 [file] [log] [blame]
# Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
Collection of decorators to help optimize Python functions and classes
by caching return values.
Return values are (by default) keyed off of the function's parameters.
Consequently, any parameters that are used in memoization must be hashable.
This library offers two styles of memoization:
1) Absolute memoization (memo) uses the full set of parameters against a
per-function memo dictionary. If this is used on an instance method, the
'self' parameter is included in the key.
2) Instance memoization (memo_i) uses a per-instance memoization dictionary.
The dictionary is stored as a member ('_memo__dict') of of the instance.
Consequently, the 'self' parameter is no longer needed/used in the
memoization key, removing the need to have the instance itself support being
Memoized function state can be cleared by calling the memoized function's
'memo_clear' method.
import inspect
class MemoizedFunction(object):
"""Handles the memoization state of a given memoized function."""
# Memoization constant to represent 'un-memoized' (b/c 'None' is a valid
# result)
class _EMPTY(object):
def __init__(self, func, ignore):
func: (function) The function to memoize
ignore: (container) The names of 'func' parameters to ignore when
generating its memo key. Only parameters that have no effect on the
output of the function should be included.
self.func = func
self.ignore = ignore
self.im_self = None
self.im_class = None
def __repr__(self):
properties = [str(self.func)]
if self.im_self is not None:
properties.append('bound=%s' % (self.im_self,))
if len(self.ignore) > 0:
properties.append('ignore=%s' % (','.join(self.ignore)))
return '%s(%s)' % (type(self).__name__, ', '.join(properties))
def __get__(self, obj, klass=None):
# Make this callable class a bindable Descriptor
if klass is None:
klass = type(obj)
self.im_self = obj
self.im_class = klass
return self
def _get_memo_dict(self):
"""Returns: (dict) the memoization dictionary to store return values in."""
raise NotImplementedError()
def _get_call_args(self, args):
"""Returns the call arguments, factoring in 'self' if this method is bound.
if self.im_self is not None:
return (self.im_self,) + args
return args
def _key(self, args, kwargs):
"""Returns: the memoization key for a set of function arguments.
This 'ignored' parameters are removed prior to generating the key.
call_params = inspect.getcallargs(self.func, *args, **kwargs)
return tuple((k, v)
for k, v in sorted(call_params.iteritems())
if k not in self.ignore)
def _call_memo(self, memo_key, args, kwargs):
"""Retrieves the memoized function result.
If the memoized function has not been memoized, it will be invoked;
otherwise, the memoized value will be returned.
memo_key: The generated memoization key for this invocation.
args, kwargs: Function parameters (only used if not memoized yet)
The memoized function's return value.
memo_dict = self._get_memo_dict()
result = memo_dict.get(memo_key, self.EMPTY)
if result is self.EMPTY:
result = memo_dict[memo_key] = self.func(*args, **kwargs)
return result
def memo_clear(self, *args, **kwargs):
"""Clears memoization results for a given set of arguments.
If no memoized result is stored for the supplied parameters, this function
is a no-op.
args, kwargs: Memoization function parameters whose memoized value should
be cleared.
memo_key = self._key(args, kwargs)
self._get_memo_dict().pop(memo_key, None)
class StaticMemoizedFunction(MemoizedFunction):
"""A MemoizedFunction bound to a static dictionary"""
def __init__(self, *args, **kwargs):
super(StaticMemoizedFunction, self).__init__(*args, **kwargs)
self.memo_dict = {}
def _get_memo_dict(self):
return self.memo_dict
def __call__(self, *args, **kwargs):
# Prepend 'self' to 'args' if bound
args = self._get_call_args(args)
# Key off full unbound/bound args
return self._call_memo(
self._key(args, kwargs),
class InstanceMemoizedFunction(MemoizedFunction):
"""A MemoizedFunction bound to an instance function's instance dictionary"""
def __init__(self, attr_name, *args, **kwargs):
super(InstanceMemoizedFunction, self).__init__(*args, **kwargs)
self.attr_name = attr_name
def _get_memo_dict(self):
# Get/create the instance-specific memo dictionary
memo_dict = getattr(self.im_self, self.attr_name, None)
if memo_dict is None:
memo_dict = {}
setattr(self.im_self, self.attr_name, memo_dict)
return memo_dict
def _key(self, args, kwargs):
# Stub out 'self' with 'None' since we are implicitly keying off of 'self'
# via its instance dictionary.
return super(InstanceMemoizedFunction, self)._key(
(None,) + args,
def __call__(self, *args, **kwargs):
# Key off bound args/kwargs (doesn't include 'self'). When we invoke the
return self._call_memo(
self._key(args, kwargs),
self._get_call_args(args), # Prepend 'self' to 'args'
def memo(ignore=()):
"""Generic function memoization decorator.
This memoizes a specific function using a function key.
The following example memoizes the function absolutely. It will be executed
once and, after that, will only returned cached results.
def my_method(print_debug_output=False):
# Perform complex calculation
if print_debug_output:
print 'The calculated result is: %r' % (result)
return result
The following example memoizes for unique values of 'param1' and 'param2',
but not 'print_debug_output' since it doesn't affect the function's result:
def my_method(param1, param2):
# Perform complex calculation using 'param1' and 'param2'
if print_debug_output:
print 'The calculated result for (%r, %r) is: %r' %
(param1, param2, result)
return result
ignore: (list) The names of parameters to ignore when memoizing.
- If None, a function-global memoization dictionary will be used
def wrap(func):
return StaticMemoizedFunction(
return wrap
def memo_i(ignore=(), attr_name=None):
"""Memoization decorator which uses an instance-local dictionary.
Instance methods can use this decorator to store memoization state in a
local instance dictionary rather than the default class-wide dictionary. This
offers several advantages:
- The instance's memoization will be collect when it is collected
- Memoization will avoid the instance's __eq__ and __hash__ functions
(allowing them to be memoized).
class MyClass(object):
def func(self):
# Perform complicated calculation
return result
ignore: (list) The names of parameters to ignore when memoizing.
attr_name: (None/str) The name of the instance-local attribute to use as the
memoization dictionary; if 'None', a default variable name will be
def wrap(func):
return InstanceMemoizedFunction(
attr_name or '_memo__dict',
return wrap