Python如何在上下文管理器__exit__中获取__enter__方法的返回值

方案1：使用独立上下文包装对象（无并发风险，推荐） #

不要直接在连接池类上实现__enter__/__exit__，每次调用时生成独立的上下文管理器实例，每个实例单独存储本次申请的连接，完全避免多线程/嵌套调用的冲突问题：

from psycopg2.pool import ThreadedConnectionPool

class ManagedConnectionPool:
    def __init__(self, *args, **kwargs):
        self._pool = ThreadedConnectionPool(*args, **kwargs)
    
    def __call__(self):
        pool = self._pool
        # 每次调用返回独立的上下文对象
        class PoolConnectionContext:
            def __enter__(self):
                self.conn = pool.getconn()
                return self.conn
            def __exit__(self, exc_type, exc_val, exc_tb):
                if self.conn:
                    # 异常自动回滚事务，避免脏连接放回池
                    if exc_type is not None:
                        self.conn.rollback()
                    pool.putconn(self.conn)
        return PoolConnectionContext()

使用方式只需加一对括号即可：

pool = ManagedConnectionPool(minconn=1, maxconn=10, dbname="your_db", user="your_user")
with pool() as conn:
    with conn.cursor() as cur:
        # 执行SQL操作
        cur.execute("SELECT 1")
        print(cur.fetchone())

方案2：使用线程本地存储（贴合你需要的语法） #

如果要严格实现with pool as conn的写法，可以借助threading.local()存储当前线程持有的连接，不同线程的存储完全隔离，不会出现连接串用的问题：

import threading
from psycopg2.pool import ThreadedConnectionPool

class ContextThreadedConnectionPool(ThreadedConnectionPool):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 线程本地存储，每个线程独立持有自己的连接引用
        self._thread_local = threading.local()
    
    def __enter__(self):
        conn = self.getconn()
        self._thread_local.current_conn = conn
        return conn
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        conn = getattr(self._thread_local, 'current_conn', None)
        if conn:
            if exc_type is not None:
                conn.rollback()
            self.putconn(conn)
            # 清理本地存储避免残留引用
            del self._thread_local.current_conn

使用方式完全符合你的预期：

pool = ContextThreadedConnectionPool(minconn=1, maxconn=10, dbname="your_db", user="your_user")
with pool as conn:
    with conn.cursor() as cur:
        # 执行SQL操作
        ...

两种方案都可以在多线程环境下安全运行，第一种方案兼容性更强，就算同一线程内同时申请多个池连接也不会冲突，更推荐使用。

内容的提问来源于stack exchange，提问作者winwin