Index and Slice a Generator in Python

If it's a 1-use slice than you could simply use the method written by ~unutbu. If you need to slice multiple times you would have to store all the intermediate values so you can "rewind" the iterator. Since iterators can iterate anything it would not have a rewind method by default.

Also, since a rewinding iterator has to store every intermediate result it would (in most cases) have no benefit over simply doing list(iterator)

Basically... you either don't need an iterator, or you're not specific enough about the situation.


import itertools

class Indexable(object):
    def __init__(self,it):
        self.it = iter(it)
    def __iter__(self):
        return self.it
    def __getitem__(self,index):
        try:
            return next(itertools.islice(self.it,index,index+1))
        except TypeError:
            return list(itertools.islice(self.it,index.start,index.stop,index.step))

You could use it like this:

it = Indexable(fib())
print(it[10])
#144
print(it[2:12:2])
#[610, 1597, 4181, 10946, 28657]

Notice that it[2:12:2] does not return [3, 8, 21, 55, 144] since the iterator had already advanced 11 elements because of the call to it[10].

Edit: If you'd like it[2:12:2] to return [3, 8, 21, 55, 144] then perhaps use this instead:

class Indexable(object):

    def __init__(self, it):
        self.it = iter(it)
        self.already_computed = []

    def __iter__(self):
        for elt in self.it:
            self.already_computed.append(elt)
            yield elt

    def __getitem__(self, index):
        try:
            max_idx = index.stop
        except AttributeError:
            max_idx = index
        n = max_idx - len(self.already_computed) + 1
        if n > 0:
            self.already_computed.extend(itertools.islice(self.it, n))
        return self.already_computed[index]

This version saves the results in self.already_computed and uses those results if possible. Otherwise, it computes more results until it has sufficiently many to return the indexed element or slice.


To slice generator you can use islice function from itertools

from itertools import islice

for i in islice(generator, 5):
    # Will be taken first 5 elems

for i in islice(generator, 5, None):
    # Will be taken everything starting at 5th