管理生成器的 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()]

这种方法将正确地迭代生成器的获取和释放，但它不会将结果传递给上下文。因此，该列表将具有值 ['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 是一个普通的 python 函数，我们可以尝试这个解决方法:

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 是任意函数而不是生成器，所以这一切都正确，但我们必须传入一个额外的语句“yield from 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!