◐ Shell
clean mode source ↗

GH-96793: Specialize FOR_ITER for generators. by markshannon · Pull Request #98772 · python/cpython

I'm seeing a 35% speedup on this benchmark:

Tree iterator
import timeit

class Tree:
    def __init__(self, left, value, right):
        self.left = left
        self.value = value
        self.right = right


    def __iter__(self):
        if self.left:
            for item in self.left:
                yield item
        yield self.value
        if self.right:
            for item in self.right:
                yield item

def tree(input: range) -> Tree | None:
    n = len(input)
    if n == 0:
        return None
    i = n // 2
    return Tree(tree(input[:i]), input[i], tree(input[i + 1:]))

def setup():
    global iterable
    assert list(tree(range(10))) == list(range(10))
    iterable = tree(range(100000))
    
print(timeit.timeit("for _ in iterable: pass", "setup()", globals=globals(), number=10))

Note that this uses yield not yield from. yield from compiles to the SEND instruction which is for another PR.


And a 36% speedup on this one:

Flat iterator
import timeit

class RangeWrapper:
    def __init__(self, n):
        self.r = range(n)

    def __iter__(self):
        for item in self.r:
            yield item

def setup():
    global iterable
    iterable = RangeWrapper(1000000)

print(timeit.timeit("for _ in iterable: pass", "setup()", globals=globals(), number=10))

Comparing the fastest of 6 runs for main and this PR.