Incorrect and overly-complicated islice example implementation in docs

[FHTMitchell]

Documentation

The itertools.islice reference implementation in the docs is complex and also slightly incorrect. For reference it is here

def islice(iterable, *args):
    s = slice(*args)
    start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
    it = iter(range(start, stop, step))
    try:
        nexti = next(it)
    except StopIteration:
        # Consume *iterable* up to the *start* position.
        for i, element in zip(range(start), iterable):
            pass
        return
    try:
        for i, element in enumerate(iterable):
            if i == nexti:
                yield element
                nexti = next(it)
    except StopIteration:
        # Consume to *stop*.
        for i, element in zip(range(i + 1, stop), iterable):
            pass

This does not correctly handle stop == 0, overwriting it to be sys.maxsize and also assumes its input is an iterator, not any iterable.

It is also just hard to reason about. Instead I propose

def islice(iterable, *args):
      s = slice(*args)
      start = s.start if s.start is not None else 0
      stop = s.stop
      step = s.step if s.step is not None else 1
      it = iter(iterable)
      for _ in zip(range(start), it):
          # Consume up to *start* position
          pass
      if stop is not None and stop <= start:
          return
      for i, element in enumerate(it, start):
          if (i - start) % step == 0:
              yield element
          if stop is not None and i + 1 >= stop:
              return

This matches islice exactly for these inputs

for start in itertools.chain(range(6), [None]):
    for stop in itertools.chain(range(6), [None]):
        for step in itertools.chain(range(1, 6), [None]):
            for r in [range(0), range(4)]:
                it = iter(r)
                myx = list(islice(it, start, stop, step))
                myy = list(it)
                it = iter(r)
                pyx = list(itertools.islice(it, start, stop, step))
                pyy = list(it)
                msg = "ERR" if myx != pyx or myy != pyy else "ok"
                print(
                    f"{msg}: {r=}, {start=}, {stop=}, {step=}, {myx=}, {pyx=}, {myy=}, {pyy=}"
                )

Linked PRs

• gh-118476: update the example implementation for islice in docs #118466
• gh-118476: Fix corner cases in islice() rough equivalent. #118559
• [3.12] gh-118476: Fix corner cases in islice() rough equivalent. (Gh-118559) #118587
• [3.11] gh-118476: Fix corner cases in islice() rough equivalent. (Gh-118559) #118588

[rhettinger]

This does not correctly handle stop == 0, overwriting it to be sys.maxsize and also assumes its input is an iterator, not any iterable.

Thanks for noticing this.

It is also just hard to reason about.

The original equivalent code was short and clear:

def islice(iterable, *args):
    s = slice(*args)
    it = iter(range(s.start or 0, s.stop or sys.maxsize, s.step or 1))
    nexti = next(it)
    for i, element in enumerate(iterable):
        if i == nexti:
            yield element
            nexti = next(it)

The current code was written by Cheryl Sabella to fix https://bugs.python.org/issue27212 . It does look a little gross but that is mostly because of the newer requirement to catch StopIteration.

Instead I propose ...

This does look nicer.

For the last loop, try to come up with a variant that doesn't use a modulo operation in an inner loop. The real C code doesn't do that. Try to mirror the C code logic as closely as possible so that the performance implications of the code are clear.

Also, the new code needs to make it clear that the iterator gets consumed all the way to the stop argument. The current code does a good job at doing that.

Besides the changes you've already looked at, see if you can find a clean way to make the equivalent not depend on both slice() and range(). Ideally, if it is not too messy, we want to communicate the actual start, stop, step, and cnt logic that is used in the C code rather than invisibly delegating it without explanation to slice() and range().

When someone looks to this rough equivalent to help their understanding, what they most likely are interested is the index computation. That is the interesting and complicated part. Once the indices are known, the actual looping and yielding are mostly obvious. The only part of the looping that needs emphasis is the final resting place for the iterator when the islice has been exhausted. This is important because use of islice() in loop relies on knowing exactly where the previous islice stopped.

[pochmann3]

Another (also passing your tests):

def islice(iterable, *args):
    s = slice(*args)
    start = s.start or 0
    stop = s.stop
    step = s.step or 1
    indices = count() if stop is None else range(max(stop, start))
    nexti = start
    for i, element in zip(indices, iterable):
        if i == nexti:
            yield element
            nexti += step

Attempt This Online!

[pochmann3]

Or with your first half including the explicit consume-to-start loop:

def islice(iterable, *args):
      s = slice(*args)
      start = s.start if s.start is not None else 0
      stop = s.stop
      step = s.step if s.step is not None else 1
      it = iter(iterable)
      for _ in zip(range(start), it):
          # Consume up to *start* position
          pass
      indices = count(start) if stop is None else range(start, stop)
      nexti = start
      for i, element in zip(indices, it):
          if i == nexti:
              yield element
              nexti += step

Attempt This Online!