memoize.py 8.59 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
#!/usr/bin/env python
#
# memoize.py - Memoization decorators.
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#
"""This module provides a handful of decorators which may be used to memoize
a function:

 .. autosummary::
    :nosignatures:

13
14
    memoize
    Memoize
15
    Instanceify
16
    memoizeMD5
17
    skipUnchanged
18
19
"""

20

21
import logging
22
import hashlib
23
import functools
Paul McCarthy's avatar
Paul McCarthy committed
24
import six
25

26
27
log = logging.getLogger(__name__)

28

29
def memoize(func=None):
30
    """Memoize the given function by the value of the input arguments.
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
    This function simply returns a :class:`Memoize` instance.
    """

    return Memoize(func)


class Memoize(object):
    """Decorator which can be used to memoize a function or method. Use like
    so::

        @memoize
        def myfunc(*a, **kwa):
            ...

        @memoize()
        def otherfunc(*a, **kwax):
            ...

    A ``Memoize`` instance maintains a cache which contains ``{args : value}``
    mappings, where ``args`` are the input arguments to the function, and
    ``value`` is the value that the function returned for those arguments.
    When a memoized function is called with arguments that are present in the
    cache, the cached values are returned, and the function itself is not
    called.


    The :meth:`invalidate` method may be used to clear the internal cache.


61
62
    Note that the arguments used for memoization must be hashable, as they are
    used as keys in a dictionary.
63
64
    """

65

66
67
68
69
70
71
72
73
74
75
    def __init__(self, *args, **kwargs):
        """Create a ``Memoize`` object.
        """

        self.__cache      = {}
        self.__func       = None
        self.__defaultKey = '_memoize_noargs_'

        self.__setFunction(*args, **kwargs)

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
    def invalidate(self, *args, **kwargs):
        """Clears the internal cache. If no arguments are given, the entire
        cache is cleared. Otherwise, only the cached value for the provided
        arguments is cleared.
        """

        if len(args) + len(kwargs) == 0:
            self.__cache = {}

        else:
            key = self.__makeKey(*args, **kwargs)

            try:
                self.__cache.pop(key)
            except KeyError:
                pass


    def __setFunction(self, *args, **kwargs):
        """Used internally to set the memoized function. """

        if self.__func is not None:
            return False

        # A no-brackets style
        # decorator was used
        isfunc = (len(kwargs) == 0 and len(args) == 1 and callable(args[0]))

        if isfunc:
            self.__func = args[0]

        return isfunc


    def __makeKey(self, *a, **kwa):
        """Constructs a key for use with the cache from the given arguments.
        """
114
        key = []
115

116
117
        if a   is not None: key += list(a)
        if kwa is not None: key += [kwa[k] for k in sorted(kwa.keys())]
118

119
120
121
122
        # This decorator was created without
        # any arguments specified - use the
        # default cache key.
        if len(key) == 0:
123
124
125
126
127
128
129
130
131
132
133
134
135
            key = [self.__defaultKey]

        return tuple(key)


    def __call__(self, *a, **kwa):
        """Checks the cache against the given arguments. If a cached value
        is present, it is returned. Otherwise the memoized function is called,
        and its value is cached and returned.
        """

        if self.__setFunction(*a, **kwa):
            return self
136

137
        key = self.__makeKey(*a, **kwa)
138

139
        try:
140
            result = self.__cache[key]
141

142
            log.debug(u'Retrieved from cache[{}]: {}'.format(key, result))
143

144
        except KeyError:
145

146
147
            result            = self.__func(*a, **kwa)
            self.__cache[key] = result
148

149
            log.debug(u'Adding to cache[{}]: {}'.format(key, result))
150

151
        return result
152
153


154
155
156
157
158
159
160
161
162
163
164
165
166
167
def memoizeMD5(func):
    """Memoize the given function. Whenever the function is called, an
    md5 digest of its arguments is calculated - if the digest has been
    previously cached, the previous value calculated by the function is
    returned.
    """

    cache = {}

    def wrapper(*args, **kwargs):
        args = list(args) + list(kwargs.values())

        hashobj = hashlib.md5()

168
169
170
171
172
        # Convert each arg to a string
        # representation, then encode
        # it into a sequence of (utf-8
        # compatible) bytes , and take
        # the hash of those bytes.
173
        for arg in args:
174
175
176
            if not isinstance(arg, six.string_types):
                arg = str(arg)
            arg = arg.encode('utf-8')
Paul McCarthy's avatar
Paul McCarthy committed
177
            hashobj.update(arg)
178
179
180
181
182
183
184
185
186

        digest = hashobj.hexdigest()
        cached = cache.get(digest)

        if cached is not None:
            return cached

        result = func(*args, **kwargs)

187
        log.debug(u'Adding to MD5 cache[{}]: {}'.format(
188
189
            digest, result))

190
191
192
193
194
        cache[digest] = result

        return result

    return wrapper
195
196
197
198


def skipUnchanged(func):
    """This decorator is intended for use with *setter* functions - a function
199
200
    which accepts a name and a value, and is intended to set some named
    attribute to the given value.
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215

    This decorator keeps a cache of name-value pairs. When the decorator is
    called with a specific name and value, the cache is checked and, if the
    given value is the same as the cached value, the decorated function is
    *not* called. If the given value is different from the cached value (or
    there is no value), the decorated function is called.

    .. note:: This decorator ignores the return value of the decorated
              function.

    :returns: ``True`` if the underlying setter function was called, ``False``
              otherwise.
    """

    import numpy as np
216

217
    cache = {}
218

219
220
221
222
223
    def wrapper(name, value, *args, **kwargs):

        oldVal = cache.get(name, None)

        if oldVal is not None:
224

225
226
227
228
            oldIsArray = isinstance(oldVal, np.ndarray)
            newIsArray = isinstance(value,  np.ndarray)
            isarray    = oldIsArray or newIsArray

229
230
231
232
233
234
235
            if isarray:
                a = np.array(oldVal, copy=False)
                b = np.array(value,  copy=False)

                nochange = (a.shape == b.shape) and np.allclose(a, b)
            else:
                nochange = oldVal == value
236
237

            if nochange:
238
                return False
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262

        func(name, value, *args, **kwargs)

        cache[name] = value

        return True

    return wrapper


class Instanceify(object):
    """This class is intended to be used to decorate other decorators, so they
    can be applied to instance methods. For example, say we have the following
    class::

        class Container(object):

            def __init__(self):
                self.__items = {}

            @skipUnchanged
            def set(self, name, value):
                self.__items[name] = value

263

264
265
266
267
268
    Given this definition, a single :func:`skipUnchanged` decorator will be
    created and shared amongst all ``Container`` instances. This is not ideal,
    as the value cache created by the :func:`skipUnchanged` decorator should
    be associated with a single ``Container`` instance.

269

270
    By redefining the ``Container`` class definition like so::
271

272

273
274
275
276
277
278
279
280
281
282
283
284
285
        class Container(object):

            def __init__(self):
                self.__items = {}

            @Instanceify(skipUnchanged)
            def set(self, name, value):
                self.__items[name] = value


    a separate :func:`skipUnchanged` decorator is created for, and associated
    with, every ``Container`` instance.

286

287
288
289
290
291
    This is achieved because an ``Instanceify`` instance is a descriptor. When
    first accessed as an instance attribute, an ``Instanceify`` instance will
    create the real decorator function, and replace itself on the instance.
    """

292

293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
    def __init__(self, realDecorator):
        """Create an ``Instanceify`` decorator.

        :arg realDecorator: A reference to the decorator that is to be
                            *instance-ified*.
        """

        self.__realDecorator = realDecorator
        self.__func          = None


    def __call__(self, func):
        """Called immediately after :meth:`__init__`, and passed the method
        that is to be decorated.
        """
        self.__func = func
        return self


    def __get__(self, instance, cls):
        """When an ``Instanceify`` instance is accessed as an attribute of
        another object, it will create the real (instance-ified) decorator,
        and replace itself on the instance with the real decorator.
        """

        if instance is None:
            return self.__func

        method    = functools.partial(self.__func, instance)
        decMethod = self.__realDecorator(method)

        setattr(instance, self.__func.__name__, decMethod)
        return functools.update_wrapper(decMethod, self.__func)