CODEWITHSUNDEEP
pythonpython-itertools
michaelyin
michaelyin

Reputation: 83

Why is the Python version of itertools.islice not working?

I have noticed that Recipes section in the itertools documentation has some useful code, like consume, but I have found a problem.

from itertools import islice
numbers = iter(range(10))
for i in numbers:
    print i 
    next(islice(numbers, 3, 3), None)

the code above works fine, it prints [0,4,8], the next(islice(numbers, 3, 3), None) trick is taken from the consume recipe.

We can see islice here is playing an important role. But, when I replace itertools.islice with the code provided in python documentation for that function, the code is not working as I expected:

#from itertools import islice

def islice(iterable, *args):
    # islice('ABCDEFG', 2) --> A B
    # islice('ABCDEFG', 2, 4) --> C D
    # islice('ABCDEFG', 2, None) --> C D E F G
    # islice('ABCDEFG', 0, None, 2) --> A C E G
    s = slice(*args)
    it = iter(xrange(s.start or 0, s.stop or 200, s.step or 1))
    nexti = next(it)
    for i, element in enumerate(iterable):
        if i == nexti:
            yield element
            nexti = next(it)


numbers = iter(range(10))
for i in numbers:
    print i 
    next(islice(numbers, 3, 3), None)

the output is different from the output above; all numbers in the range are printed and the islice() consume line doesn't appear to do anything.

Can anybody can explain why this is and the way islice works here?

Upvotes: 4

Views: 3262

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1122232

The slice start and stop arguments are equal, so the next(it) call on the xrange raises StopIteration:

>>> s = slice(3, 3, None)
>>> it = iter(xrange(s.start or 0, s.stop or 200, s.step or 1))
>>> next(it)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

and the slice() generator function immediately terminates. This is only logical, the xrange() function is not going to produce any numbers in that case.

The itertools documentation is illustrative, and not meant to be complete. You've found one edgecase where the sample Python code does not work the same way that the C implementation works.

To make the Python code work for this case as well, catch the StopIteration exception, and test for the edgecase explicitly in the loop:

def islice(iterable, *args):
    s = slice(*args)
    it = iter(xrange(s.start or 0, s.stop or 200, s.step or 1))
    try:
        nexti = next(it)
        consume = False
    except StopIteration:
        if s.start < 1:
            return
        nexti = s.start - 1
        consume = True
    for i, element in enumerate(iterable):
        if i == nexti:
            if consume:
                return
            yield element
            nexti = next(it)

The Python sample code uses a slice() object to make the code more compact and easier to read; the C code does its own start, stop and step parsing, and handles the count itself instead of delegating to xrange(), and thus doesn't run into this edgecase.

Upvotes: 4

Related Questions