Pythonic way to manage generators

Question

Does Python offer some syntactic sugar to sweeten up the construction of a generator test as below?

def acquire():
    print('Acquiring resource')
    yield 'A' 
    
def do_stuff():
    print('Doing stuff')
    yield 'B'
    
def release():
    print('Releasing resource')
    yield 'C'

def test():
    yield from acquire()
    yield from do_stuff()
    yield from release()
    
[u for u in test()] # has value ['A', 'B', 'C']

Basically I want a syntax that allows acuire and release to appear in the same statement. At first I thought a context manager would be appropriate, for instance:

class conman:
    def __init__(self, acq, rel):
        self.acq = acq
        self.rel = rel
        
    def __enter__(self):
        try:
            while True:
                next(self.acq)
        except StopIteration:
            return
    
    def __exit__(self, _, __, ___):
        try:
            while True:
                next(self.rel)
        except StopIteration:
            return

def conmantest():
    with conman(acquire(), release()):
        yield from do_stuff()
[u for u in conmantest()]

This approach will iterate correctly through the generators acquire and release, but it doesn't pass the result on to the conmantext. As a result, the list will have value ['B'], even though it still prints all the messages in the right order.

A different approach is to use a decorator

def manager(acq, rel):
    def decorator(func):
        def wrapper(*args, **kwargs):
            yield from acq
            yield from func(*args, **kwargs)
            yield from rel
            return
        return wrapper
    return decorator

@manager(acquire(), release())
def do_stuff_decorated():
    print('Doing stuff')
    yield 'B'

[u for u in do_stuff_decorated()]

This does the right thing, but in practice do_stuff is a list of statements and it is not always desirable to write a generator around them.

If release were an ordinary pythonfunction we could try this work around:

class conman2:
    def __init__(self, acq, rel):
        self.acq = acq
        self.rel = rel
        
    def __enter__(self):
        return self.acq
    
    def __exit__(self, _, __, ___):
        self.rel()
        
def release_func():
    print('Releasing stuff')

def conman2test():
    with conman2(acquire(), release_func) as r:
        yield from r
        yield from do_stuff()
[u for u in conmantest()]

This does all the right things, since release_func is an arbitrary function and not a generator, but we had to pass in an extra statement `yield from r'. Something like this is used in the SimPy library for discrete event programming to implement a context for resources where they are automatically released when the context ends.

However, I was hoping there might be some syntax like

class yielded_conman:
    def __init__(self, acq, rel):
        self.acq = acq
        self.rel = rel
        
    def __yielded_enter__(self):
        yield from self.acq()

    def __yielded_exit__(self, _, __, ___):
        yield from self.rel()

def yieldconmantest():
    with yielded_conman(acquire(), release()):
        yield from do_stuff()

[u for u in conmantest()] # has value ['A', 'B', 'C']

which does all the right things.

Answer 1

One approach using contextlib:

from contextlib import contextmanager

def acquire():
    print('Acquiring resource')
    yield 'A'

def do_stuff():
    print('Doing stuff')
    yield 'B1'
    raise Exception('Something happened!')
    yield 'B2'

def release():
    print('Releasing resource')
    yield 'C'

@contextmanager
def cntx(a, b, c):
    def _fn():
        try:
            yield from a
            yield from b
        finally:
            yield from c

    try:
        yield _fn
    finally:
        pass

def fn():
    with cntx(acquire(), do_stuff(), release()) as o:
        yield from o()

[print(i) for i in fn()]

Print:

Acquiring resource
A
Doing stuff
B1
Releasing resource
C
Traceback (most recent call last):
  File "main.py", line 35, in <module>
    [print(i) for i in fn()]
  File "main.py", line 35, in <listcomp>
    [print(i) for i in fn()]
  File "main.py", line 33, in fn
    yield from o()
  File "main.py", line 22, in _fn
    yield from b
  File "main.py", line 10, in do_stuff
    raise Exception('Something happened!')
Exception: Something happened!