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