Index and Slice a Generator in Python

Lets say I have a generator function that looks like this:

def fib():
    x,y = 1,1
    while True:
        x, y = y, x+y
        yield x

Ideally, I could just use fib()[10] or fib()[2:12:2] to get indexes and slices, but currently I have to use itertools for these things. I can't use a generator for a drop in replacement for lists.

I believe the solution will be to wrap fib() in a class:

class Indexable(object):
    ....

fib_seq = Indexable(fib())

What should Indexable look like to make this work?

Solution

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.