I have a generator function A.
For example (in reality I have a more complex function A),
def A():
yield from [i**2 for i in range(20)]
Writing another generator function B, I want to enumerate all elements that A returns except the first element.
What are concise ways to implement this in Python 3?
Use itertools.islice:
itertools.islice(generator,1,None)
Usually, you don't need this in an expression, so you just call next(it), ignoring the results, to consume and discard the first element.
However, if the iterator might be empty, you have to decide what you want to happen:
StopIteration, in which case next(it) is fine.next(it) inside an except StopIteration: raise SomethingElse().next(it, None). You can find examples of these in the stdlib and docs. For example, if you scan through the recipes in itertools:
def pairwise(iterable):
"s -> (s0,s1), (s1,s2), (s2, s3), ..."
a, b = tee(iterable)
next(b, None)
return zip(a, b)
This is doing exactly what you want to do—skip the first element of b. And if iterable is empty, you don't want an error here; you just want to iterate nothing. So, next(b, None).
What if you need to do this in the middle of an expression?
Then you can write a function that skips the first element:
def skip_first(iterable):
it = iter(iterable)
next(it, None)
return it
(Again, you have to decide what you want to happen for an empty iterable.)
This returns a first-skipped version of the iterator, so you can use it inline. (It also mutates the iterator you passed in, of course, but you normally only use on a temporary value that you're not keeping any references to, so that's not a problem.)
Or, if you need to return a generator instead of an arbitrary iterator (usually you don't):
def skip_first(iterable):
it = iter(iterable)
next(it, None)
yield from it
Or you can use the more general version of the same idea, itertools.islice. The following have the same effect:
it = skip_first(it)
it = itertools.islice(it, 1, None)
While we're on the itertools recipes, it's worth looking at consume:
def consume(iterator, n=None):
"Advance the iterator n-steps ahead. If n is None, consume entirely."
# Use functions that consume iterators at C speed.
if n is None:
# feed the entire iterator into a zero-length deque
collections.deque(iterator, maxlen=0)
else:
# advance to the empty slice starting at position n
next(islice(iterator, n, n), None)
Forget the None part; the interesting bit is that it skips n elements with an islice and a next. (Notice that it's mutating iterator in-place, not returning something.)
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With