blob: 3eed568112b9bfe1cb30aa9ae2375328efc7de62 [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
# Instance variable added to class instances that use memoization to
# differentiate their memoization values.
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=None, memo_dict=None):
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 = frozenset(ignore or ())
self.memo_dict = memo_dict
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(sorted(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_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 _get_memo_dict(self):
"""Returns: (dict) the memoization dictionary to store return values in."""
memo_dict = None
if self.im_self is not None:
# Is the instance dictionary defined?
memo_dict = getattr(self.im_self, MEMO_INSTANCE_VARIABLE, None)
if memo_dict is None:
memo_dict = {}
setattr(self.im_self, MEMO_INSTANCE_VARIABLE, memo_dict)
return memo_dict
# No instance dict; use our local 'memo_dict'.
if self.memo_dict is None:
self.memo_dict = {}
return self.memo_dict
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.
if self.im_self is not None:
# We are bound to an instance; use None for args[0] ("self").
args = (None,) + tuple(args)
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__(self, *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()
memo_key = self._key(args, kwargs)
result = memo_dict.get(memo_key, self.EMPTY)
if result is self.EMPTY:
args = self._get_call_args(args)
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 arguments are supplied, this will clear all retained memoization for
this function.
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_dict = self._get_memo_dict()
if args or kwargs:
memo_key = self._key(args, kwargs)
memo_dict.pop(memo_key, None)
def memo(ignore=None, memo_dict=None):
"""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.
def wrap(func):
return MemoizedFunction(
return wrap