scripts/common/memo.py

# 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
   hashed.

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):
    pass
  EMPTY = _EMPTY()

  def __init__(self, func, ignore):
    """
    Args:
      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.

    Args:
      memo_key: The generated memoization key for this invocation.
      args, kwargs: Function parameters (only used if not memoized yet)
    Returns:
      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:
      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),
        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,
        kwargs,
    )

  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'
        kwargs,
    )


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.

  @memo.memo(ignore=('print_debug_output',))
  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:

  @memo.memo()
  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

  Args:
    ignore: (list) The names of parameters to ignore when memoizing.
    memo_dict:
        - If None, a function-global memoization dictionary will be used
  """
  def wrap(func):
    return StaticMemoizedFunction(
        func,
        ignore=ignore,
    )
  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):
    @memo.memo_i()
    def func(self):
      # Perform complicated calculation
      return result

  Args:
    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
        used.
  """
  def wrap(func):
    return InstanceMemoizedFunction(
        attr_name or '_memo__dict',
        func,
        ignore=ignore,
    )
  return wrap