A decorator for generators that makes them indexable like a sequence.
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 | from itertools import count
class indexediterator(object):
# Helper class to be returned as the actual generator with
# indexing; wraps the generator in an iterator that also
# supports item retrieval by index.
def __init__(self, gen):
self.__gen = gen # the generator that created us
self.__iter = gen._iterable()
def __iter__(self):
# Return the generator function; note that we return
# the same one each time, which matches the semantics
# of actual generators (i.e., once the generator function
# is called, iter(gen) returns the same iterator and does
# not reset the state)
return self.__iter
def next(self):
# Return next item from generator
term = next(self.__iter)
if term == self.__gen.sentinel:
raise StopIteration
return term
def __len__(self):
# If the generator is exhausted, we know its length, so
# we can use that information; if not, we raise TypeError,
# just like any other object with no length
result = self.__gen._itemcount()
if result is None:
raise TypeError, "object of type %s has no len()" % \
self.__class__.__name__
return result
def __getitem__(self, index):
# Return the item at index, advancing the generator if
# necessary; if the generator is exhausted before index,
# raise IndexError, just like any other sequence when an
# index out of range is requested
result = self.__gen._retrieve(index)
if result is self.__gen.sentinel:
raise IndexError, "sequence index out of range"
return result
class IndexedGenerator(object):
"""Make a generator indexable like a sequence.
"""
sentinel = object()
def __init__(self, gen):
# The underlying generator
self.__gen = gen
# Memoization fields
self.__cache = []
self.__iter = None
self.__empty = False
def _retrieve(self, n):
# Retrieve the nth item from the generator, advancing
# it if necessary
# Negative indexes are invalid unless the generator
# is exhausted, so check that first
if n < 0:
end = self._itemcount()
if (end is None) or (end == 0):
# No length known yet, or no items at all
return self.sentinel
else:
return self.__cache[end + n]
# Now try to advance the generator (which may empty it,
# or it may already be empty)
while (not self.__empty) and (n >= len(self.__cache)):
try:
term = next(self.__iter)
except StopIteration:
self.__empty = True
else:
self.__cache.append(term)
if n < len(self.__cache):
return self.__cache[n]
return self.sentinel
def _iterable(self):
# Yield terms from the generator
for n in count():
term = self._retrieve(n)
if term is self.sentinel:
break
yield term
def _itemcount(self):
# Once we are exhausted, the number of items in the
# sequence is known, so we can provide it; otherwise
# we return None
if self.__empty:
return len(self.__cache)
return None
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)
return indexediterator(self)
# 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/
indexable_generator = partial(DelayedDecorator, IndexedGenerator)
|
This is an enhanced version of the MemoizeGenerator
decorator recipe at http://code.activestate.com/recipes/577992-memoize-generator/, which makes the generator indexable like a sequence as well. Note that the sequence is not mutable. Also note that, while it behaves like a sequence for indexing, it behaves like a generator with regard to iteration; repeated calls to next
(or something like a for
loop) will eventually exhaust it and StopIteration
will be raised, even though indexing into it will still retrieve items already yielded.
This decorator shares the same general limitations as MemoizedGenerator
. It can only be used directly to decorate an ordinary function, and it is not thread-safe; it assumes that all realizations of the 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, to use this decorator on a method of a class, you will need the Delayed Decorator recipe as well.
Finally, this class, like MemoizedGenerator
, caches 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. This implementation also supports slicing, which the above recipe does not (since it adds considerable complication).