Python3 从一个函数返回多个上下文管理器，以在单个 with 语句中使用

Question

给定：

con = psycopg2.connect()

with con, con.cursor() as c:
  c.execute() # some query inside here

根据 psycopg2 文档https://www.psycopg.org/docs/usage.html#transactions-control，con 对象管理事务并负责数据库事务的提交和回滚。因此，

con

语句中需要

con.cursor()

和

with

才能正确管理提交/回滚

现在我想多次重复

with

部分代码，做多个交易，比如

con = psycopg2.connect()

with con, con.cursor() as c:
  c.execute() # some query inside here

with con, con.cursor() as c:
  c.execute() # another query inside here

...

with con, con.cursor() as c:
  c.execute() # final query inside here

这行得通，但这需要我为每个

con, con.cursor()

块复制粘贴

with

语句的

with

部分。

现在我想知道是否有可能在 python 中创建一个函数来返回我可以直接传递给

with

语句的函数，以将

con, con.cursor()

减少到

some_custom_function()

沿着这些线的东西：

con = psycopg2.connect()

def cursor():
  return con, con.cursor() # this doesn't work

with cursor() as c:
  c.execute() # some query inside here

with cursor() as c:
  c.execute() # another query inside here

...

with cursor() as c:
  c.execute() # final query inside here

（您可能想知道为什么，但是

con.cursor()

方法也接受诸如

cursor_factory=psycopg2.extras.RealDictCursor

之类的参数。然后我将不得不在每个

with

语句中重复这些参数。但是为了这个例子的简单性，我已经把它排除在外。）

Answer 1

尝试

contextlib.contextmanager

：

from contextlib import contextmanager

import psycopg2

con = psycopg2.connect(...)


@contextmanager
def with_txn_and_cursor():
    with con, con.cursor(cursor_factory=RealDictCursor) as cur:
        yield cur


with with_txn_and_cursor() as cur:
    cur.execute(...)

如果您同时需要

con

和

cur

，请从上下文管理器中生成一个元组。

@contextmanager
def with_txn_and_cursor_2():
    with con, con.cursor(cursor_factory=RealDictCursor) as cur:
        yield (con, cur)


with with_txn_and_cursor_2() as (con, cur):
    cur.execute(...)