HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonMinor

Universal memoization decorator

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
decoratoruniversalmemoization

Problem

I've just written a simple caching / memoization python decorator. It's purpose is to cache what the function returns for all the arguments combinations it's been ever invoked with.

So, if, say, we had a function my_function(a, b=2, c=3) invoked once with my_function(1, 2, c=3), the results should be then cached whenever I call it with my_function(1, 2, 3) or my_function(1, c=3, b=2) or my_function(c=3, b=2, a=1).

Is what's below a good approach and could it be somehow improved?

```
import functools
import pickle

def memoize_all(func):
"""
This is a caching decorator. It caches the function results for
all the arguments combinations, so use it with care. It does not matter whether the arguments
are passed as keywords or not.
"""
cache = {}

@functools.wraps(func)
def cached(*args, **kwargs):
arg_names = func.func_code.co_varnames
arg_dict = {}
for i, arg in enumerate(args):
arg_dict[arg_names[i]] = args[i]
arg_dict.update(**kwargs)
key = pickle.dumps(arg_dict)
if key not in cache:

Solution

I don't know what's up with the spacing in the docstring, but fix it.

Your memoize_all doesn't work with everything:

memoize_all(int)(1)
#>>> Traceback (most recent call last):
#>>>   ...
#>>> AttributeError: type object 'int' has no attribute 'func_code'


I suppose this doesn't bother you much, but it irks me ;).

Consider

for i, arg in enumerate(args):
    arg_dict[arg_names[i]] = args[i]


There's no point doing args[i]; it's just arg. Further, one should do

for i, (name, arg) in enumerate(zip(arg_names, args)):
    arg_dict[name] = arg


But then one can just do

arg_dict = dict(zip(arg_names, args))


You could even use

arg_dict = dict(zip(arg_names, args), **kwargs)


You try and save this with

key = pickle.dumps(arg_dict)


I'm not at all fond of this method. pickle is a bit fragile and more complicated than I would expect. I would personally do something like

key = frozenset(arg_dict.items())


If you skip generating arg_dict, this is

key = frozenset(kwargs.items()) | frozenset(zip(arg_names, args))


This does prevent accepting mutable arguments, but it's normally best to enforce purity before caching arguments anyway.

You then have just

def cached(*args, **kwargs):
    arg_names = func.func_code.co_varnames
    key = frozenset(kwargs.items()) | frozenset(zip(arg_names, args))

    if key not in cache:
        cache[key] = func(*args, **kwargs)
    return cache[key]


Note that in Python 3, this would be better with a signature object:

import functools
from inspect import signature

def memoize_all(func):
    """
    This is a caching decorator. It caches the function results for
    all the arguments combinations, so use it with care. It does not
    matter whether the arguments are passed as keywords or not.
    """
    cache = {}
    func_sig = signature(func)

    @functools.wraps(func)
    def cached(*args, **kwargs):
        bound = func_sig.bind(*args, **kwargs)
        key = frozenset(bound.arguments.items())

        if key not in cache:
            cache[key] = func(*args, **kwargs)
        return cache[key]
    return cached

Code Snippets

memoize_all(int)(1)
#>>> Traceback (most recent call last):
#>>>   ...
#>>> AttributeError: type object 'int' has no attribute 'func_code'
for i, arg in enumerate(args):
    arg_dict[arg_names[i]] = args[i]
for i, (name, arg) in enumerate(zip(arg_names, args)):
    arg_dict[name] = arg
arg_dict = dict(zip(arg_names, args))
arg_dict = dict(zip(arg_names, args), **kwargs)

Context

StackExchange Code Review Q#78371, answer score: 6

Revisions (0)

No revisions yet.