Using yield to run code after method of function execution

November 26, 2022

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?

>Solution :

You need another next call to have it run the code after the yield:

database = Database()
gen = database()  # Saved the generator to a variable
conn = next(gen)
cur = conn.cursor()
cur.execute("select * from users")
result = cur.fetchall()
conn.commit()
next(gen)  # Triggers the latter part of the function

Also note, when you exhaust a generator, it raises a StopIteration exception as you’ll see. You’ll need to catch that as well.