Index and Slice a Generator in Python

Lets say I have a generator function that looks like开发者_JAVA技巧 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?

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

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.

So based off the code from ~unutbu and adding in a little itertools.tee:

import itertools

class Indexable(object):
    def __init__(self, it):
        self.it = it

    def __iter__(self):
        self.it, cpy = itertools.tee(self.it)
        return cpy

    def __getitem__(self, index):
        self.it, cpy = itertools.tee(self.it)
        if type(index) is slice:
            return list(itertools.islice(cpy, index.start, index.stop, index.step))
        else:
            return next(itertools.islice(cpy, index, index+1))

For small cases (e.g. you are only ever going to use at most the first 20 Fiboonacci numbers), where it isn't worth adding a full class, the simplest solution is to convert islice output to a list:

from itertools import islice

list(islice(fib(), 20))

Here is ~unutbu's answer modified to subclass list. Obviously abuse such as append, insert etc. will produce weird results!

you do get __str__ and __repr__ methods for free though

import itertools
class Indexable(list):
    def __init__(self,it):
        self.it=it
    def __iter__(self):
        for elt in self.it:
            yield elt
    def __getitem__(self,index):
        try:
            max_idx=index.stop
        except AttributeError:
            max_idx=index
        while max_idx>=len(self):
            self.append(next(self.it))
        return list.__getitem__(self,index)