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

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.

Answer

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")