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