获取__enter__的返回值(以__exit__为单位)

Question

我试图使数据库连接池的行为像一个简单的连接(psycopg2)。就像这样:

from psycopg2.pool import ThreadedConnectionPool    
pool = ThreadedConnectionPool(...)
with pool as conn:
with conn.cursor() as cur:
# do sql stuff
...

要做到这一点，我们应该从池中conn = pool.getconn()，并在我们完成查询后，pool.putconn(conn)它回来。这似乎是一个使用上下文管理器的完美示例。

然而，如果我们这样做:

# this function -> a part of ThreadedConnectionPool subclass
def __enter__(self):
return self._pool.getconn()  # doesn't return self, returns a database connection

则需要在__exit__方法中返回连接。比如:

# this function -> a part of ThreadedConnectionPool subclass
def __exit__(self, *args):
self._pool.putconn(conn)  # <--- here!
...

但是我们如何访问conn对象呢?似乎上下文应该知道它，但是在哪里可以访问它

?我读过这样的问题,但它在多线程环境中使用这个对象同时连接和商店很多,所以没有机会我们可以使用一个本地状态函数,是吗?

如何在__exit__函数调用中访问__enter__返回值?

Answer 1

此特定对象一次只能有一个连接。所以这应该可以工作:

def __enter__(self):
self.conn = self.__pool.getconn()
return self.conn
def __exit__(self, *args):
self.__pool.putconn(self.conn)
self.conn = None