Day 19 · ~19m

Python Context Managers and the with Statement: Guaranteed Cleanup

Priya's connection closes at the end of her script—unless an exception fires. With the with statement, cleanup is absolutely guaranteed, no matter what happens.

student (thinking)

Okay so I have this database connection at the top of my script, and I close it at the bottom. Works fine.

teacher (neutral)
connection = get_db_connection()
try:
    result = connection.execute(query)
    process(result)
finally:
    connection.close()
teacher (serious)

That finally block is you fighting the language. You know connection.close() needs to run, so you wrapped it in finally to make sure. But you had to know to do that. And you had to write it every single time.

student (curious)

So there's a better way?

teacher (neutral)

The with statement. Watch:

with get_db_connection() as connection:
    result = connection.execute(query)
    process(result)

No try, no finally, no explicit close(). When the with block ends — even if an exception fires — connection.close() runs automatically.

student (confused)

How does it know to call close()? It's not in my code.

teacher (focused)

Because get_db_connection() returns something called a context manager. It's an object with two special methods: __enter__ and __exit__. When you use with, Python calls __enter__ when you step into the block, and __exit__ when you leave it. No matter how you leave.

student (thinking)

So the object itself is responsible for cleanup?

teacher (encouraging)

Exactly. The object says: "I know how to set myself up and how to tear myself down. You don't have to remember." It's like a fire exit. You don't have to carry the fire extinguisher out of the building every time you leave. The exit is always there, always available.

student (curious)

Show me what it looks like inside.

teacher (neutral)

Here's a simple context manager:

class Database:
    def __init__(self, connection_string):
        self.connection_string = connection_string
        self.conn = None
    
    def __enter__(self):
        print(f"Opening connection to {self.connection_string}")
        self.conn = really_open_connection(self.connection_string)
        return self.conn  # This is what 'as connection' gets
    
    def __exit__(self, exc_type, exc_value, traceback):
        print("Closing connection")
        if self.conn:
            self.conn.close()

Now you can write:

with Database("postgres://...") as connection:
    result = connection.execute(query)
student (confused)

What are exc_type, exc_value, traceback in exit? Those are weird names.

teacher (serious)

They're the exception information. If an exception fires inside the with block, exit receives the details. If no exception fires, all three are None.

student (thinking)

So exit always runs, but it gets information about whether something went wrong?

teacher (focused)

Yes. And if exit returns True, it suppresses the exception. If it returns False (or returns nothing, which is the default), the exception keeps propagating.

student (curious)

So the cleanup code can decide whether to hide an error?

teacher (serious)

It can. But usually it doesn't. Usually exit cleans up and lets the exception bubble up so the caller knows something went wrong.

def __exit__(self, exc_type, exc_value, traceback):
    print("Closing connection")
    if self.conn:
        self.conn.close()
    return False  # Don't suppress exceptions
student (thinking)

So the pattern is: enter sets things up and returns the resource. exit tears things down, no matter what.

teacher (encouraging)

That's the whole thing.

student (focused)

But writing a whole class just to close a connection seems like overkill.

teacher (neutral)

That's why Python has contextlib.contextmanager. It's a decorator that turns a function with yield into a context manager.

from contextlib import contextmanager

@contextmanager
def database_connection(connection_string):
    print(f"Opening {connection_string}")
    conn = really_open_connection(connection_string)
    try:
        yield conn  # This is what 'as connection' gets
    finally:
        print("Closing")
        conn.close()

# Now you use it exactly like before:
with database_connection("postgres://...") as connection:
    result = connection.execute(query)
student (surprised)

Wait, that's cleaner. You just yield the resource, then put cleanup code after the yield?

teacher (proud)

And if an exception fires during the with block, the finally still runs. Yield is the boundary. Code before yield is enter. Code after yield is exit.

student (thinking)

So @contextmanager is just syntactic sugar for writing enter and exit manually?

teacher (focused)

Yes. Pick whichever is clearer. A class when you have multiple methods and state to track. The decorator when it's simple: setup, yield, cleanup.

student (curious)

Can I see something more realistic? Like, something from your codebase?

teacher (neutral)

Sure. Imagine you're processing orders and you want to track whether the whole operation succeeded before you commit to the database.

@contextmanager
def transaction(database):
    print("Starting transaction")
    database.begin()
    try:
        yield database
    except Exception as e:
        print(f"Rolling back due to {e}")
        database.rollback()
        raise  # Re-raise the exception
    else:
        print("Committing")
        database.commit()

Now when you process orders, the context manager handles rollback and commit for you:

def process_orders(orders, database):
    with transaction(database) as db:
        for order in orders:
            db.insert_order(order)  # If this fails, rollback happens automatically
student (excited)

So if one order fails, the entire transaction rolls back? And you don't have to write the rollback code every time?

teacher (proud)

Exactly. The context manager is responsible for transaction lifecycle. Your code just says what to do, and the context manager handles success and failure paths.

student (thinking)

But what if I need to do something after the context manager closes? Like, log a message after the connection closes?

teacher (serious)

You can't easily do that inside exit. exit runs as the context manager closes, so at that point your resource is already cleaned up. If you need post-cleanup code, you'd handle it outside the with block.

with database_connection("postgres://...") as connection:
    result = connection.execute(query)

print("Connection is closed now")  # This runs after the context manager exits
student (focused)

Okay so with statement: setup, use, cleanup. All automatic. And contextmanager decorator is the easiest way to write simple ones.

teacher (neutral)

For most cases, yes. Use the class when you need enter and exit to be separate methods — like when you have other methods on the context manager object itself. Use @contextmanager for the simple yield-based pattern.

student (curious)

When would I actually need to write context managers? Like, in what situations?

teacher (focused)

Anywhere you have setup and teardown. Database connections — set up connection, execute queries, close connection. File handles — Python's open() is a context manager, which is why you write with open() as f:. Locks in multithreaded code — acquire lock, use resource, release lock. Database transactions — begin, execute, commit or rollback.

student (thinking)

So with open() is using the context manager protocol?

teacher (encouraging)

Yes. File objects implement enter and exit. enter makes sure the file is open. exit closes it. That's why people say "always use with when you open files" — the context manager guarantees closure.

student (focused)

And if you don't use with, you have to remember to close manually?

teacher (serious)

Or the file stays open until Python garbage collects the object. Which might be soon, might be much later. With is the guarantee. No ambiguity.

student (curious)

Can you nest context managers? Like, open two files at once?

teacher (neutral)

Yes:

with open("input.txt") as input_file:
    with open("output.txt", "w") as output_file:
        output_file.write(input_file.read())

Or more concisely:

with open("input.txt") as input_file, open("output.txt", "w") as output_file:
    output_file.write(input_file.read())

Both input_file and output_file go out of scope at the end, and both close automatically.

student (excited)

That's clean. Both resources get cleaned up in the right order.

teacher (amused)

And if an exception fires, both still close. Try doing that with manual close() calls without a try/finally for each one.

student (focused)

Alright so here's my question: when I see with in Amir's code, what's actually happening?

teacher (neutral)

Amir is using an object that implements enter and exit. When execution reaches the with statement, Python calls enter. When execution leaves the with block for any reason — normal exit, exception, break, return — exit runs.

student (thinking)

So with is not language-level magic. It's a protocol. Objects opt into the protocol by implementing enter and exit.

teacher (proud)

Correct. And you can write context managers for anything that needs setup and teardown. Want a context manager that times a block of code? Done.

import time
from contextlib import contextmanager

@contextmanager
def timer(label):
    start = time.time()
    try:
        yield
    finally:
        elapsed = time.time() - start
        print(f"{label}: {elapsed:.2f}s")

with timer("database query"):
    perform_expensive_query()
student (excited)

So you're saying I could write a context manager that logs when I enter and exit a function? For debugging?

teacher (encouraging)

Absolutely. Context managers are tools for any kind of setup and cleanup. Files, connections, transactions, timing, logging, state management. Anywhere you have a "before this block" and "after this block" responsibility.

student (focused)

Okay I get it. But here's where I'm confused. Yesterday we learned about exception handlers — try/except/finally. Today we're learning with statements. Are they solving the same problem?

teacher (serious)

Different problems. Try/except handles what to do when an error occurs. With handles setup and teardown. They work together. With usually has a try/finally under the hood, but you don't write it.

student (thinking)

So the context manager is the try/finally, and you're just not seeing it.

teacher (focused)

Yes. The context manager author writes the try/finally (or the enter/exit protocol, which Python translates to try/finally). You just write with and trust that cleanup happens.

student (curious)

What's the order? Does exit run before or after my except block?

teacher (neutral)

exit is part of the with protocol. If an exception fires inside the with block, exit runs first. Then the exception propagates outside the with. If your code has an except outside the with, that catches the propagated exception.

try:
    with database_connection() as conn:
        conn.execute(bad_query)  # Raises ValueError
except ValueError:
    print("Query was invalid")

Order: execute bad_query → ValueError raised → exit runs (cleanup) → ValueError caught by except.

student (surprised)

So exit runs before my except?

teacher (focused)

exit is called automatically. Your except is outside the with block, so it catches whatever exception exit doesn't suppress.

student (thinking)

So if I'm using with statements for setup and cleanup, I don't need to write finally blocks myself?

teacher (encouraging)

Correct. The context manager handles it. That's the whole point. You describe what you want to do (use a database, open a file, lock a resource), and the context manager handles the guaranteed cleanup.

student (focused)

Tomorrow we're learning what?

teacher (amused)

Exception hierarchies and chains. What happens when one exception triggers another, and how to keep track of the whole story.

student (excited)

So tomorrow it gets even more complex?

teacher (serious)

Yes. Today: "This thing needs cleanup, guaranteed." Tomorrow: "What if cleanup fails? What if five things fail in sequence and you need to know which one caused which?" Today you're guaranteeing exit. Tomorrow we make sure that exit knows what went wrong.