Pythonic式的生成器管理方式

Question

Python是否提供了一些语法糖来美化生成器test的构造，如下所示？

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']

基本上，我想要一个允许acuire和release出现在同一语句中的语法。起初，我认为上下文管理器是合适的，例如：

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()]

这种方法将正确地迭代生成器获取和释放，但它不会将结果传递给conmantext。因此，该列表将具有值['B']，即使它仍然以正确的顺序打印所有消息。

另一种方法是使用装饰器。

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()]

这样做是正确的，但在实践中，do_stuff是一个语句列表，并不总是希望围绕它们编写一个生成器。

如果release是一个普通的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()]

这做了所有正确的事情，因为release_func是一个任意函数，而不是一个生成器，但我们必须传入一个额外的语句‘来自r’。在离散事件编程的SimPy library中使用类似这样的东西来实现资源的上下文，当上下文结束时，资源会自动释放。

但是，我希望有一些语法，比如

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']

它做了所有正确的事情。

Answer 1

一种方法是使用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()]

打印：

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!