← Home

116. Contextlib Module

Context managers with @contextmanager, suppress, and closing

116. Contextlib Module

πŸ”’ A resource manager uses context managers for guaranteed cleanup. The developer wrote a class-based context manager but forgot exit, and wrote a generator-based context manager but forgot yield β€” both silently break cleanup logic.

πŸ’‘ Fun fact: The with statement and context managers were added in Python 2.5 (2006) via PEP 343, written by Guido van Rossum. Before that, developers had to use try/finally blocks to guarantee cleanup. The contextlib module was added at the same time to make writing context managers easier. Python’s built-in open() has been a context manager since Python 2.6 β€” with open("file.txt") as f: is the idiomatic Python way to open files because it guarantees f.close() even if an exception occurs inside the block.

⚠️ Watch out: A class-based context manager without __exit__ does not raise an error β€” Python falls back to the default object.__exit__, which silently does nothing. Your cleanup code never runs, the connection never closes, and resources leak. Similarly, a @contextmanager function without yield raises RuntimeError: generator didn't yield at the with statement β€” Python requires exactly one yield to mark the boundary between setup and teardown code.

πŸ€” Think about it: __exit__(self, exc_type, exc_val, exc_tb) receives the exception details if one occurred in the with block. If you return True from __exit__, the exception is suppressed β€” execution continues after the with block as if no exception happened. When would suppressing an exception in __exit__ be the right design choice? What could go wrong if you accidentally return True?

Learning objectives

  • Implement enter and exit for class-based context managers
  • Use @contextmanager with yield to create generator-based context managers
  • Understand that yield in a @contextmanager marks the body of the with block
  • Use contextlib.suppress to silently ignore specific exception types
  • Return False from exit to let exceptions propagate normally

Key concepts

  • enter(self) β€” setup, returns value for β€˜as’ clause
  • exit(self, exc_type, exc_val, exc_tb) β€” teardown
  • @contextmanager + yield β€” generator-based context manager
  • contextlib.suppress(ExcType) β€” ignore specific exceptions
  • finally in @contextmanager β€” guaranteed cleanup

Try it

Concept detail

contextlib β€” Context Managers Made Easy

Context managers run setup and teardown code reliably, even when exceptions occur. The with statement uses them.

Class-Based Context Manager

class ManagedFile:
    def __init__(self, path):
        self.path = path
        self.file = None

    def __enter__(self):
        self.file = open(self.path, "r")
        return self.file  # assigned to "as" variable

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.file:
            self.file.close()
        return False  # False = don't suppress exceptions

@contextmanager Decorator

from contextlib import contextmanager

@contextmanager
def managed_connection(db_path):
    conn = sqlite3.connect(db_path)
    try:
        yield conn       # <-- body of with block runs here
    finally:
        conn.close()     # always runs, even on exception

The yield is mandatory and marks the boundary of the with block.

contextlib.suppress

from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("maybe_exists.txt")
# If file doesn't exist, FileNotFoundError is silently ignored

contextlib.closing

from contextlib import closing
import urllib.request

with closing(urllib.request.urlopen("http://example.com")) as page:
    content = page.read()
# page.close() is called automatically

The exit Signature

def __exit__(self, exc_type, exc_val, exc_tb):
    # exc_type, exc_val, exc_tb are None if no exception occurred
    self.cleanup()
    return False  # return True to suppress the exception

Solution

import contextlib
import sqlite3

class ManagedConnection:
    """A context manager that opens a SQLite connection and closes it after use."""

    def __init__(self, db_path=":memory:"):
        self.db_path = db_path
        self.conn = None
        self.closed = False

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_path)
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        # FIX 1: __exit__ closes the connection and marks cleanup as done
        if self.conn:
            self.conn.close()
        self.closed = True
        return False  # don't suppress exceptions

@contextlib.contextmanager
def temporary_table(conn, table_name):
    """Create a temp table, yield the connection, then drop the table."""
    conn.execute(f"CREATE TABLE IF NOT EXISTS {table_name} (id INTEGER, value TEXT)")
    try:
        yield conn  # FIX 2: yield marks the body of the with block
    finally:
        conn.execute(f"DROP TABLE IF EXISTS {table_name}")

def suppress_key_error(func, *args, **kwargs):
    """Run func(*args) and suppress KeyError if it occurs. Return None on error."""
    with contextlib.suppress(KeyError):
        return func(*args, **kwargs)
    return None

Tests

import sqlite3

def test_managed_connection_provides_connection():
    """__enter__ should return a usable sqlite3 connection."""
    mgr = ManagedConnection()
    with mgr as conn:
        assert conn is not None
        # Can execute SQL β€” confirms it's a real connection
        conn.execute("CREATE TABLE t (x INTEGER)")

def test_managed_connection_calls_exit():
    """__exit__ must be called β€” tracked via self.closed flag."""
    mgr = ManagedConnection()
    with mgr as conn:
        pass  # just enter and exit
    assert mgr.closed is True, (
        "__exit__ was never called β€” did you define __exit__ on ManagedConnection?"
    )

def test_managed_connection_closes_on_exception():
    """__exit__ should be called even when an exception occurs in the with block."""
    mgr = ManagedConnection()
    try:
        with mgr as conn:
            raise ValueError("intentional error")
    except ValueError:
        pass
    assert mgr.closed is True, (
        "__exit__ must be called even when an exception occurs in the with block"
    )

def test_temporary_table_yields():
    """@contextmanager function must have a yield."""
    conn = sqlite3.connect(":memory:")
    try:
        with temporary_table(conn, "test_temp") as c:
            assert c is not None
            c.execute("INSERT INTO test_temp VALUES (1, 'hello')")
            cursor = c.execute("SELECT * FROM test_temp")
            rows = cursor.fetchall()
            assert len(rows) == 1
    except RuntimeError as e:
        assert False, f"temporary_table raised RuntimeError β€” is yield missing? {e}"
    conn.close()

def test_temporary_table_drops_after_with_block():
    """Table should be dropped after the with block exits."""
    conn = sqlite3.connect(":memory:")
    with temporary_table(conn, "temp_data") as c:
        c.execute("INSERT INTO temp_data VALUES (42, 'test')")
    # After with block, table should be gone
    try:
        conn.execute("SELECT * FROM temp_data")
        table_exists = True
    except Exception:
        table_exists = False
    assert not table_exists, "temporary_table should DROP the table after the with block"
    conn.close()

def test_suppress_key_error_on_missing_key():
    """suppress_key_error should return None when KeyError is raised."""
    d = {"a": 1}
    result = suppress_key_error(lambda: d["missing_key"])
    assert result is None

def test_suppress_key_error_returns_value_on_success():
    """suppress_key_error should return the actual value when no error."""
    d = {"x": 42}
    result = suppress_key_error(lambda: d["x"])
    assert result == 42

def test_suppress_does_not_suppress_other_errors():
    """suppress(KeyError) should NOT suppress TypeError or ValueError."""
    import contextlib
    def raise_type_error():
        raise TypeError("not a KeyError")
    caught = False
    try:
        suppress_key_error(raise_type_error)
    except TypeError:
        caught = True
    assert caught, "suppress(KeyError) should not swallow TypeError"

Resources