Using yield to run code after method of function execution

Question

I'm trying to create a class method that can run some code after its execution.

In pytest we have this functionality with fixtures:

@pytest.fixture
def db_connection(conn_str: str):
    connection = psycopg2.connect(conn_str)
    yield connection
    connection.close() # this code will be executed after the test is done

Using this fixture in some test guarantees that connection will be closed soon after the test finishes. This behavior is described here, in the Teardown section.

When I try to do it in my own class methods, I didn't get the same result.

class Database:
    def __call__(self, conn_str: str):
        conn = psycopg2.connect(conn_str)
        yield conn
        print("Got here")
        conn.close()

database = Database()
conn = next(database())
cur = conn.cursor()
cur.execute("select * from users")
result = cur.fetchall()
conn.commit()

result

The output is the data in users table, but I never see the "Got here" string, so I'm guessing this code after the yield keyword never runs.

Is there a way to achieve this?

Answer 1

What you are trying to do is implement a context manager; the similarly to a Pytext fixture is incidental.

You can do this with contextmanager.contextlib

from contextlib import contextmanager

@contextmanager
def db_connection(conn_str):
    connection = psycopg2.connect(conn_str)
    yield connection
    connection.close()

with db_connection(...) as db:
    ...

or define Database.__enter__ and Database.__exit__ explicitly:

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

    def __enter__(self):
        self.conn = psycopg2.connect(self.conn_str)
        return self.conn

    def __exit__(self, *args):
        print("Got here")
        self.conn.close()

with Database(...) as db:
    ...

(You can use the connection returned by psycopg2.connect as a context manager itself.)