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

A decorator for generators that makes them indexable like a sequence.

Python, 121 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
 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).

Created by Peter Donis on Sun, 1 Jan 2012 (MIT)
Python recipes (4591)
Peter Donis's recipes (4)

Required Modules

Other Information and Tasks