Use Python generator's .send() when generator is wrapped to act as a context manager

user_6582418 0
Viewed 165

Python's contextlib provides wrappers to turn generators into context managers:

from contextlib import contextmanager

@contextmanager
def gen():
    yield

with gen() as cm:
    ...

And generators provide the ability to send values into generators that just yielded:

def gen():
    x = yield
    print(x)

g = gen()
next(g)
g.send(1234) # prints 1234 and raises a StopIteration because we have only 1 yield

Is there any way to get both behaviors at the same time? I would like to send a value into my context manager so that it can be used while handling __exit__. So something like this:

from contextlib import contextmanager

@contextmanager
def gen():
    x = yield
    # do something with x

with gen() as cm:
    # generator already yielded from __enter__, so we can send
    something.send(1234)

I'm not sure if this is a good/reasonable idea or not. I feel like it does break some layer of abstraction since I would be assuming that the context manager was implemented as a wrapped generator.

If it is a viable idea, I'm not sure what something should be.

generator python
2 Answers

The generator underlying a @contextmanager is directly accessible via its gen attribute. Since the generator cannot access the context manager, the latter must be stored before the context:

from contextlib import contextmanager

@contextmanager
def gen():
    print((yield))  # first yield to suspend and receive value...
    yield           # ... final yield to return at end of context

manager = gen()     # manager must be stored to keep it accessible
with manager as value:
    manager.gen.send(12)

It is important that the generator has exactly the right amount of yield points - @contextmanager ensures that the generator is exhausted after exiting the context.

@contextmanager will .throw raised exceptions in the context, and .send None when done, for which the underlying generator can listen:

@contextmanager
def gen():
    # stop when we receive None on __exit__
    while (value := (yield)) is not None:
        print(value)

In many cases, it is probably easier to implement the context manager as a custom class, though. This avoids complications from using the same channel to send/recv values and pause/resume the context.

class SendContext:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def send(self, value):
        print("We got", value, "!!!")


with SendContext() as sc:
    sc.send("Deathstar")
user_53499160

I don't think that is possible, and there is probably a better way to implement what you want.

The way that contextmanager is implemented is described in PEP 344 about the with statement. The contextmanager calls the next method twice on the user's generator function, once in the __enter__ method and once in the __exit__ method, both without sending in any value.

The way you might implement what you want would be roughly (untested):

def create_context_manager(value_you_want_to_pass):
    @contextmanager
    def gen():
        do_something_with(value_you_want_to_pass)
        yield the_object_for_with_statement
        do_something_else_with(value_you_want_to_pass)
    return gen

with create_context_manager(1234) as cm:
    do_something_with(cm)
user_26472790
Related