如何在自定义类方法中利用yield实现类似pytest fixture的后置执行逻辑？

你遇到的核心问题是：生成器函数在调用next()后只会执行到yield处暂停，只有当生成器被完全耗尽（或显式关闭）时，yield之后的代码才会执行。你当前只调用了一次next()，生成器处于暂停状态，所以后续的打印和连接关闭代码根本没机会运行。

下面提供两种可靠的解决方案：

方案一：用上下文管理器（推荐） #

直接把Database类实现为上下文管理器，通过__enter__和__exit__方法分别处理资源初始化和清理逻辑，这是Python中管理资源生命周期的标准方式，和pytest fixture的效果一致。

import psycopg2

class Database:
    def __init__(self, conn_str: str):
        self.conn_str = conn_str
        self.conn = None

    def __enter__(self):
        # 初始化资源：建立数据库连接
        self.conn = psycopg2.connect(self.conn_str)
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 清理资源：无论是否发生异常，都会执行这里的代码
        print("Got here")
        if self.conn:
            self.conn.close()

# 使用方式：with语句会自动调用__enter__和__exit__
with Database("你的数据库连接字符串") as conn:
    cur = conn.cursor()
    cur.execute("select * from users")
    result = cur.fetchall()
    conn.commit()

print(result)

方案二：正确处理生成器生命周期 #

如果坚持用生成器实现，必须确保生成器被完全遍历或显式触发后续代码。可以用try-finally包裹生成器逻辑，再通过遍历生成器来触发清理：

import psycopg2

class Database:
    def __call__(self, conn_str: str):
        conn = psycopg2.connect(conn_str)
        try:
            yield conn
        finally:
            # 无论生成器是否正常结束，finally块都会执行
            print("Got here")
            conn.close()

# 使用方式：用for循环遍历生成器（会自动耗尽生成器）
database = Database()
for conn in database("你的数据库连接字符串"):
    cur = conn.cursor()
    cur.execute("select * from users")
    result = cur.fetchall()
    conn.commit()

print(result)

补充说明 #

pytest的fixture之所以能自动执行yield后的代码，是因为pytest框架内部会帮你管理生成器的生命周期：在测试用例执行完毕后，会自动继续迭代生成器，直到触发StopIteration，从而执行yield后的清理逻辑。而你手动调用生成器时，必须自己完成这个过程。

内容的提问来源于stack exchange，提问作者André Carvalho