在函数执行方法之后使用 yield 运行代码答案

【问题标题】：Using yield to run code after method of function execution在函数执行方法之后使用 yield 运行代码
【发布时间】：2022-11-27 11:03:48
【问题描述】：

我正在尝试创建一个可以在执行后运行一些代码的类方法。

在 pytest 中，我们通过 fixtures 拥有此功能：

@pytest.fixture
def db_connection(conn_str: str):
    connection = psycopg2.connect(conn_str)
    yield connection
    connection.close() # this code will be executed after the test is done

在某些测试中使用此夹具可确保连接将在测试完成后立即关闭。在拆解部分中描述了此行为 here。

当我尝试在我自己的类方法中这样做时，我没有得到相同的结果。

class Database:
    def __call__(self, conn_str: str):
        conn = psycopg2.connect(conn_str)
        yield conn
        print("Got here")
        conn.close()

database = Database()
conn = next(database())
cur = conn.cursor()
cur.execute("select * from users")
result = cur.fetchall()
conn.commit()

result

输出是 users 表中的数据，但我从未看到“Got here”字符串，所以我猜测 yield 关键字从未运行后的这段代码。

有没有办法做到这一点？

【问题讨论】：

标签： python oop pytest generator psycopg2

【解决方案1】：

你需要另一个 next 调用让它在 yield 之后运行代码：

database = Database()
gen = database()  # Saved the generator to a variable
conn = next(gen)
cur = conn.cursor()
cur.execute("select * from users")
result = cur.fetchall()
conn.commit()
next(gen)  # Triggers the latter part of the function

另请注意，当您耗尽生成器时，它会引发 StopIteration 异常，如您所见。你也需要抓住它。

【讨论】：