Can Python context managers be used with generators?

Question

I'm trying to use a context manager which controls visiting a subdirectory, and it seems like it would be very elegant to combine that with a generator expression, but it doesn't seem to work. Is there any way to correct this so I can use the two together?

Here's the example:

import os, sys
from contextlib import contextmanager
from glob import glob
 
@contextmanager
def visitDir(d):
    os.chdir(d)
    yield d
    os.chdir("..")
 
paths = [os.path.join('.', p[0:-1]) for p in glob('*/')]
 
def clean():
    for p in (visitDir(p) for p in paths): # This is the magic line
    print(p)
    print(os.getcwd())
 
clean() # Context manager apparently expires within the generator expression

Answer 1

You need to control the entering and leaving of a context. The generator expression has no concept of a wider context, so you cannot just put a context manager in a generator expression and expect that to be automatically entered when yielded.

Only the with statement manages the actual context, triggering the __enter__ and __exit__ hooks on the CM. You can just use the produced context manager objects from the generator expression here:

def clean():
    for cm in (visitDir(p) for p in paths):
        with cm as p:
            print p
            print os.getcwd()

The with statement calls cm.__enter__ here, and when the block is complete cm.__exit__ is called.

But I'd find the following more readable and comprehensible:

def clean():
    for p in paths:
        with visitDir(p):
            print p
            print os.getcwd()

because creating the context manager as part of the with line is just easier to follow.