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).
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.