Python3 return multiple contextmanagers from a function to be used in a single with statement

Given:

con = psycopg2.connect()

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

According to the psycopg2 documentation https://www.psycopg.org/docs/usage.html#transactions-control, the con object manages the transaction and takes care of commit and rollback of the db transaction. So both the con and con.cursor() are required in the with statement to properly manage commit/rollback

Now I want repeat the with part of the code multiple times, to do multiple transactions, such as

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

This works but this requires me to copy paste the con, con.cursor() part of the with statement for every with block.

Now I was wondering if it is possible in python to create a function that returns something that I can pass directly to the with statement to reduce con, con.cursor() to some_custom_function()

Something along these lines:

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

(You may be wondering why, but the con.cursor() method also takes arguments such as cursor_factory=psycopg2.extras.RealDictCursor. Then I would have to repeat those arguments with every with statement as well. But for simplicity of this example, I've left that out of the question.)

Solution

Try 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(...)

If you need both the con and cur, yield a tuple out of the context manager.

@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(...)