How are context managers in unconsumed generators finalized?

Viewed 737

I do not understand how and when a context manager in a non-finished generator is closed. Consider the following context manager and function:

from contextlib import contextmanager

@contextmanager
def ctx():
    print('enter ctx')
    yield
    print('exit ctx')

def gen_nums(n):
    with ctx():
        yield from range(n)

My first intuition was that if I call gen_nums but do not consume the generator fully, ctx will never be closed, which was rather concerning. For example:

for i, j in zip(range(5), gen_nums(10)):
    print(f'{i}, {j}')

Here exit ctx is not printed at the end. As I saw it, that meant that if I had a file context in the generator it would be left open; however I then realized that doing the same with files would actually close the file properly. After some tests, I found out that if I did:

from contextlib import contextmanager

@contextmanager
def ctx():
    print('enter ctx')
    try:
        yield
    finally:
        print('exit ctx')

Now exit ctx was printed at the end. So I suppose some exception would be triggered at some point, but I don't know which, where or when (I tried to print the exception with except BaseException as e but it did not work). It seems it happens when the generator is deleted, because if I do:

g = gen_nums(10)
for i, j in zip(range(5), g):
    print(f'{i}, {j}')
del g

Then exit ctx only happens after del g. However, I would like to have a better understanding of what is happening here and who is triggering what.

1 Answers

Consider this:

@contextmanager
def ctx():
    print('enter ctx')
    try:
        print('ctx begins yield')
        yield
        print('ctx finishes yield')
    finally:    
        print('exit ctx')

def gen_nums(n):
    print('entering generator')
    l = range(n)
    with ctx():
        print('inside context manager')
        for i in l:
            print('gen before yield')
            yield i
            print('gen after yield')
        print('existing ctx')
    print('exiting generator')

The result is thus:

>>> g = gen_nums(3)
>>> next(g)
entering generator
enter ctx
ctx begins yield
inside context manager
gen before yield
0
>>> next(g)
gen after yield
gen before yield
1
>>> next(g)
gen after yield
gen before yield
2
>>> next(g)
gen after yield
exiting ctx
ctx finishes yield
exit ctx
exiting generator
Traceback (most recent call last):
  File "<pyshell#165>", line 1, in <module>
    next(g)
StopIteration

It seems like the contextmanager is only ever truly exited at the point after the generator's last iteration and before it hits StopIteration. So at the very last iteration of the generator the contextmanager would remain open if it wasn't ran in a for loop (which handles the StopIteration).

You wouldn't be able to catch the StopIteration within the contextmanager either because it happens after the generator is exhausted, at which point the contextmanager would have been exited already.

Also, IIRC, in constructing a contextmanager you should always utilize try... finally... anyways.

Related