Memoized decorators for python

memoize.py

def memoized(func):
    """
    @see http://en.wikipedia.org/wiki/Memoization
    @see http://wiki.python.org/moin/PythonDecoratorLibrary#Memoize

    Decorator used for deterministic functions: functions which will always
    return the same result for the same arguments. Functions which are
    expensive (perhaps involving database queries, read and parse files,
    etc.) can be memoized: cached into memory, so that the result can
    simply be retrieved if the function is called again with the same
    arguments.

    One of the most common examples of memoization is with the Fibonacci
    sequence, which goes like: 1, 1, 2, 3, 5, 8, 13, 21, ... The first two
    items are '1', and the subsequent items are the sum of the previous two
    items. In Python, a function which will generate the nth item in the
    sequence:

    >>> def fib(n):
    ...    if n < 2: return 1
    ...    return fib(n-2) + fib(n-1)
    ...
    >>> import timeit
    >>> timeit.Timer(lambda: fib(35)).timeit(1)
    7.63999966

    requires significant processor time (7.6s) because the function requires
    significant recursion. However, by inspection, the function is
    deterministic: the 5th item in the sequence should always be the exact
    same number, so it doesn't need to be recalculated for every sequence
    item after the 5th. With a @memoized decorator we have:

    >>> @memoized
    ... def fib(n):
    ...    if n < 2: return 1
    ...    return fib(n-2) + fib(n-1)
    ...
    >>> timeit.Timer(lambda: fib(35)).timeit(1)
    0.00021589

    Memoizing reduces the processor time down to 0.2ms. Now you can
    calculate very large values for instance, calculating the 900th item in
    the Fibonacci sequence is completely infeasible with the non-memoized
    version, but takes about 4ms with the memozied version:

    >>> fib(900)
    88793027306605937532517516910637647045239090036365766884466525589158360259770006891772711976920559280382807770394537471560061517120086971996377683290300054868066659454250625417891167369401L

    As a caveat, memoized functions cannot receive keyword arguments. Since
    dicts in Python are not hashable, they cannot be used in memoize cache
    for a memoized-decorated function.
    """
    func.cache = {}
    def decorator(*args):
        try:
            return func.cache[args]
        except KeyError:
            value = func(*args)
            func.cache[args] = value
            return value
        except TypeError:
            # uncachable -- for instance, passing a list as an
            # argument.
            # Better to not cache than to blow up entirely.
            return func(*args)
    # Transfer some information from the decorated function
    decorator.__name__ = func.__name__
    decorator.__doc__ = func.__doc__

    def invalidate(key=None):
        # Invalidates the function call cache
        if key:
            del func.cache[key]
        else:
            func.cache = {}
    decorator.invalidate = invalidate

    return decorator

def memoized_by(arg_mutator):
    """
    Slightly more complicated approach to memoizing. It allows the decorator
    to specify a function to mutate the arguments used in the cache storage
    and lookup. This is useful, for example, for a function which receives a
    date as an argument but will return the same result for any date in an
    entire month. Simple memoizing would fail because two dates in the same
    month are not equal. Therefore:

    @memoized_by(lambda m: (m.year, m.month))
    def memoized_func(month):
        return something_expensive_by_month(month)

    would cache the result of the memoized_func by the month and year of the
    received argument. Therefore, future calls to the function with dates in
    the same month would simply retrieve the same value.
    """
    def memoized(func):
        func.cache = {}
        def decorator(*args):
            # Mutate the args
            cache_args = arg_mutator(*args)
            try:
                return func.cache[cache_args]
            except KeyError:
                value = func(*args)
                func.cache[cache_args] = value
                return value
            except TypeError:
                # uncachable -- for instance, passing a list
                # as an argument.
                # Better to not cache than to blow up entirely.
                return func(*args)
        # Transfer some information from the decorated function
        decorator.__name__ = func.__name__
        decorator.__doc__ = func.__doc__

        def invalidate(key=None):
            # Invalidates the function call cache
            if key:
                del func.cache[key]
            else:
                func.cache = {}
        decorator.invalidate = invalidate

        return decorator
    return memoized