pythongeneratorcontextmanagergenerator-expression

Can Python context managers be used with generators?


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

Solution

  • 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.