Why generators don't raise StopIteration exception when called through pytest?

Question

Consider the generator:

def sample():
    print("Setup")
    yield
    print("Teardown")


gen = sample()
next(gen)
next(gen)

When I call the first next(gen), the generator executes until print("Setup") and the next time I call it, it executes until after print("Teardown"). But also, since there is no second yield statement, it raises StopIteration exception. But when I do the following modification:

import pytest
@pytest.fixture
def sample():
    print("Setup")
    yield
    print("Teardown")


def test_case(sample):
    print("Executing test case")

Now, when I run the pytest command, all the code executes properly, first the setup part, then the test_case and then the teardown part. I am sure pytest is calling the generator twice during execution since all the print statements are being executed. But how is it that the StopIteration exception not being raised here? My assumption is that pytest is handling it internally. Confirm my assumption and please correct me if I am wrong. Thank you.

Answer 1

Look at the source code to @pytest.fixture in fixtures.py.

@pytest.fixture is a decorator, so it creates a new function that looks exactly like sample(), but that has a wrapper around it. The first time the sample() is called, the wrapper expects the function to return normally. The second time sample() is called, it is expecting the function to return StopIteration, and explicitly discards it.

In general, you should not expect a function that has a decorator to behave the same way as that same function without a decorator.

Answer 2

next() is low-level and StopIteration is low-level too. Those are parts of the iteration protocol. The other side must obey the rules too, a StopIteration must be caught (or it becomes a RuntimeError - PEP479).

The StopIteration exception is always raised at the end of the generator and if you don't see an error, the caller must have caught the exception directly (try-except) or indirectly, e.g.:

gen = sample()
# the for statement speaks the iteration protocol
# behind the scenes
for v in gen:
    print(f"got {v}")