Welcome, guest | Sign In | My Account | Store | Cart

A wrapper class for generators that "memoizes" them, so that even if the generator is realized multiple times, each term only gets computed once (after that the result is simply returned from a cache).

Python, 50 lines
 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
from functools import partial
from itertools import count

class MemoizedGenerator(object):
    """Memoize a generator to avoid computing any term more than once.
    """
    
    def __init__(self, gen):
        # The underlying generator
        self.__gen = gen
        # Memoization fields
        self.__cache = []
        self.__iter = None
        self.__empty = False
    
    def __call__(self, *args, **kwargs):
        """Make instances of this class callable.
        
        This method must be present, and must be a generator
        function, so that class instances work the same as their
        underlying generators.
        """
        if not (self.__empty or self.__iter):
            self.__iter = self.__gen(*args, **kwargs)
        for n in count():
            # First check the cache
            if n < len(self.__cache):
                yield self.__cache[n]
            # See if another copy of the generator emptied it
            # since our last iteration
            elif self.__empty:
                break
            # If none of the above, advance the generator
            # (which may empty it)
            else:
                try:
                    term = next(self.__iter)
                except StopIteration:
                    self.__empty = True
                    break
                else:
                    self.__cache.append(term)
                    yield term


# This creates a decorator that works if applied to a method
# (the above will only work on an ordinary generator function)
# -- requires the Delayed Decorator recipe at
# http://code.activestate.com/recipes/577993-delayed-decorator/
memoize_generator = partial(DelayedDecorator, MemoizedGenerator)

Generic implementation of a "memoized" generator for cases where each term of the generator is expensive to compute, but it is known that every realization of the generator will compute the same terms. This class allows multiple realizations of the generator to share a single computation of each term.

The MemoizeGenerator class can be used to wrap a generator directly, but it only works for ordinary functions (i.e., not methods). For ease of use and flexibility, it is recommended that the memoize_generator decorator be used instead, since that automatically handles both ordinary functions and methods.

Note that this recipe is not thread-safe; it assumes that all realizations of the memoized generator run in the same thread, so that it is guaranteed that no more than one realization will be mutating the memoization fields at a time.

Also note that to use this decorator on a method of a class, you will need the Delayed Decorator recipe as well.

Finally, note that this class memoizes all realizations of its underlying generator, even if they are invoked with different arguments. This is probably not the desired behavior for generators that take arguments, since different arguments may correspond to different terms being yielded. A more "production" version of this decorator that handles this issue, by memoizing separately for each distinct set of arguments, is available in the plib library at http://pypi.python.org/pypi/plib in the Python Package Index.