Index and Slice a Generator in Python

Question

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?

Answer 1

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.