-
Notifications
You must be signed in to change notification settings - Fork 220
/
memoization.py
126 lines (100 loc) · 3.74 KB
/
memoization.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
from collections.abc import Hashable
from functools import partial
from itertools import tee
__all__ = ['memoized_func', 'memoized_meth', 'memoized_generator']
class memoized_func(object):
"""
Decorator. Caches a function's return value each time it is called.
If called later with the same arguments, the cached value is returned
(not reevaluated). This decorator may also be used on class methods,
but it will cache at the class level; to cache at the instance level,
use ``memoized_meth``.
Adapted from: ::
https://wiki.python.org/moin/PythonDecoratorLibrary#Memoize
"""
def __init__(self, func):
self.func = func
self.cache = {}
def __call__(self, *args):
if not isinstance(args, Hashable):
# Uncacheable, a list, for instance.
# Better to not cache than blow up.
return self.func(*args)
if args in self.cache:
return self.cache[args]
else:
value = self.func(*args)
self.cache[args] = value
return value
def __repr__(self):
"""Return the function's docstring."""
return self.func.__doc__
def __get__(self, obj, objtype):
"""Support instance methods."""
return partial(self.__call__, obj)
class memoized_meth(object):
"""
Decorator. Cache the return value of a class method.
Unlike ``memoized_func``, the return value of a given method invocation
will be cached on the instance whose method was invoked. All arguments
passed to a method decorated with memoize must be hashable.
If a memoized method is invoked directly on its class the result will not
be cached. Instead the method will be invoked like a static method: ::
class Obj(object):
@memoize
def add_to(self, arg):
return self + arg
Obj.add_to(1) # not enough arguments
Obj.add_to(1, 2) # returns 3, result is not cached
Adapted from: ::
code.activestate.com/recipes/577452-a-memoize-decorator-for-instance-methods/
"""
def __init__(self, func):
self.func = func
def __get__(self, obj, objtype=None):
if obj is None:
return self.func
return partial(self, obj)
def __call__(self, *args, **kw):
if not isinstance(args, Hashable):
# Uncacheable, a list, for instance.
# Better to not cache than blow up.
return self.func(*args)
obj = args[0]
try:
cache = obj.__cache_meth
except AttributeError:
cache = obj.__cache_meth = {}
key = (self.func, args[1:], frozenset(kw.items()))
try:
res = cache[key]
except KeyError:
res = cache[key] = self.func(*args, **kw)
return res
class memoized_generator(object):
"""
Decorator. Cache the return value of an instance generator method.
"""
def __init__(self, func):
self.func = func
def __repr__(self):
"""Return the function's docstring."""
return self.func.__doc__
def __get__(self, obj, objtype=None):
if obj is None:
return self.func
return partial(self, obj)
def __call__(self, *args, **kwargs):
if not isinstance(args, Hashable):
# Uncacheable, a list, for instance.
# Better to not cache than blow up.
return self.func(*args)
obj = args[0]
try:
cache = obj.__cache_gen
except AttributeError:
cache = obj.__cache_gen = {}
key = (self.func, args[1:], frozenset(kwargs.items()))
it = cache[key] if key in cache else self.func(*args, **kwargs)
cache[key], result = tee(it)
return result