Context managers: writing your own `with`

Part of the free LogicWiz Python course, Module 14: Exceptions, context managers and testing. The examples and practice run in the app.

Meher's sales file came out empty. The program crashed halfway, before reaching f.close(), so nothing was saved. with open(...) would have saved her. What is with actually doing?

You'll learn to

  • Explain what `with` does in terms of `__enter__` and `__exit__`
  • Write a context manager as a class and with `@contextmanager`
  • Decide whether an exception passes through `__exit__` or is suppressed
  • Use `contextlib.suppress` and know where the standard library uses `with`

Learn

What `with` does

with EXPR as name:
    body
  1. Python evaluates EXPR to get a context manager object.
  2. It calls its __enter__(); whatever that returns is bound to name.
  3. It runs the body.
  4. It always calls __exit__(...), whether the body finished, returned, or raised.

So with is a tidy try/finally: set up, do the work, and the cleanup is guaranteed. Files, locks, database transactions and timers all use it.

A class-based context manager

class Shift:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print("open counter:", self.name)
        return self

    def __exit__(self, exc_type, exc, tb):
        print("close counter:", self.name)
        return False

with Shift("morning") as s:
    print("selling as", s.name)
print("after the with")

Exceptions pass through `__exit__`

__exit__(self, exc_type, exc, tb) gets three arguments:

  • no error: None, None, None
  • an error: its class, the exception object and the traceback

Then its return value decides:

  • falsy (False, None, no return): the exception keeps going up to the caller
  • truthy (True): the exception is suppressed, and the program continues after the with

Suppress only the specific errors you truly handle. A __exit__ that always returns True hides every bug inside the block.

Handling one kind of error in `__exit__`

class IgnoreBadQty:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        if exc_type is ValueError:
            print("handled:", exc)
            return True
        return False

with IgnoreBadQty():
    qty = int("two")
    print("never printed")
print("program continues")

The easy way: `@contextmanager`

from contextlib import contextmanager

@contextmanager
def till(name):
    print("open", name)
    cash = {"total": 0}
    try:
        yield cash
    finally:
        print("close", name, "with Rs", cash["total"])

with till("till-1") as t:
    t["total"] += 50
    t["total"] += 35

Code before yield is __enter__; the yielded value becomes the as name; code after it is __exit__. If the body raises, the exception appears at the yield, so the try/finally is what guarantees the cleanup.

`suppress` and several managers at once

from contextlib import suppress
import os

with suppress(FileNotFoundError):
    os.remove("yesterday.csv")
print("no crash, even though the file was missing")

with open("sales.txt", "w") as out, open("log.txt", "w") as log:
    out.write("chai,10\n")
    log.write("saved 1 line\n")
print(out.closed, log.closed)

suppress(X) is a ready-made manager whose __exit__ returns True for X only. One with can open several managers; they close in reverse order.

In a @contextmanager function, put yield inside try/finally. Without it, an error in the body jumps out at the yield and the cleanup code after it never runs, which is exactly the bug with was meant to prevent.

Where you meet with: open() files, threading.Lock(), a sqlite3 connection (commits on success, rolls back on error), decimal.localcontext(), unittest.mock.patch(...), contextlib.redirect_stdout, and pytest.raises(...). Interviewers love "write a context manager that times a block" or "that rolls back on error".

In one line

with calls __enter__ (its return is the as name) and always __exit__; a truthy __exit__ suppresses the error; with @contextmanager, wrap yield in try/finally.

Interview corner

What is a context manager? Write one.

An object with __enter__ and __exit__ used by a with statement; __exit__ is guaranteed to run, so it is used for cleanup (closing files, releasing locks, rolling back transactions).

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self
    def __exit__(self, *exc):
        self.elapsed = time.perf_counter() - self.start
        return False

What happens when an exception is raised inside a `with` block?

The rest of the block is skipped and __exit__(exc_type, exc, tb) is called with the exception's details. If __exit__ returns a truthy value the exception is suppressed; otherwise it propagates after cleanup.

How does `contextlib.contextmanager` work?

It turns a generator function into a context manager. Code before yield runs on enter, the yielded value is bound by as, and code after yield runs on exit. An exception in the block is re-raised at the yield, so you wrap it in try/finally to guarantee cleanup.