在函数执行方法后使用yield来运行代码

Question

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

在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

在某些测试中使用此夹具可以保证在测试结束后很快关闭连接。此行为在这里的Teardown部分中进行了描述。

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

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关键字之后的代码永远不会运行。

有办法做到这一点吗?

Answer 1

你要做的是实现一个上下文管理器;与Pytext fixture的类似是偶然的。

可以使用contextmanager.contextlib

from contextlib import contextmanager
@contextmanager
def db_connection(conn_str):
connection = psycopg2.connect(conn_str)
yield connection
connection.close()
with db_connection(...) as db:
...

或显式定义Database.__enter__和Database.__exit__:

class Database:
def __init__(self, conn_str: str):
self.conn_str = conn_str
def __enter__(self):
self.conn = psycopg2.connect(self.conn_str)
return self.conn
def __exit__(self, *args):
print("Got here")
self.conn.close()
with Database(...) as db:
...

(您可以使用psycopg2.connect返回的连接作为上下文管理器本身)

Answer 2

您需要另一个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异常。你也需要抓住这个