Python: Zero to Hero
Home/Functional Programming and Advanced Functions
Share

Chapter 24: Context Managers

You've been writing with open("file.txt") as f: since Chapter 8. It closes the file automatically — even if your code crashes. That's a context manager at work.

Context managers handle setup and teardown automatically. They guarantee that cleanup always happens, no matter what. Files get closed. Database connections get released. Locks get unlocked. Timers get stopped. You never forget, because you can't.

This chapter shows you exactly how they work, how to write your own using both classes and generators, and the patterns that make them so useful in real code.

The Problem They Solve

Without a context manager, you have to remember to clean up:

f = open("data.txt")
data = f.read()
# ... do something with data ...
f.close()   # easy to forget, or skip if an exception is raised

If an exception happens between open() and close(), the file never closes. Your program leaks a file handle. Do this enough times and the OS refuses to open any more files.

The try/finally pattern fixes it:

f = open("data.txt")
try:
    data = f.read()
    # ... work with data ...
finally:
    f.close()   # always runs, even if an exception occurs

This is correct — but verbose and easy to get wrong. The with statement is cleaner:

with open("data.txt") as f:
    data = f.read()
    # ... work with data ...
# file is always closed here, guaranteed

Same behavior. Zero extra code. The context manager handles the try/finally for you.

How with Works

The with statement calls two methods on the context manager object:

  1. __enter__() — runs at the start of the with block. The return value is bound to the as variable.
  2. __exit__(exc_type, exc_val, exc_tb) — runs when the block ends, whether normally or due to an exception.
with expression as variable:
    body

Python translates this into roughly:

_cm = expression
variable = _cm.__enter__()
try:
    body
except:
    if not _cm.__exit__(*sys.exc_info()):
        raise
else:
    _cm.__exit__(None, None, None)

If __exit__ returns True, the exception is suppressed. If it returns False (or None), the exception propagates.

Writing a Context Manager Class

Implement __enter__ and __exit__:

class ManagedFile:
    """Context manager for opening files safely."""
    def __init__(self, path, mode="r"):
        self.path = path
        self.mode = mode
        self.file = None

    def __enter__(self):
        print(f"Opening {self.path!r}")
        self.file = open(self.path, self.mode)
        return self.file   # bound to the 'as' variable

    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"Closing {self.path!r}")
        if self.file:
            self.file.close()
        if exc_type:
            print(f"  Exception suppressed: {exc_val}")
        return False   # don't suppress exceptions


with ManagedFile("data.txt", "w") as f:
    f.write("Hello from a context manager!\n")

# Prints:
# Opening 'data.txt'
# Closing 'data.txt'

Handling exceptions in __exit__

__exit__ receives three arguments: the exception type, value, and traceback. If no exception occurred, all three are None.

class SuppressErrors:
    """Silently ignore specific exception types."""
    def __init__(self, *exceptions):
        self.exceptions = exceptions

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type and issubclass(exc_type, self.exceptions):
            print(f"Suppressed: {exc_val}")
            return True   # swallow the exception
        return False      # let other exceptions propagate


with SuppressErrors(FileNotFoundError, PermissionError):
    open("does_not_exist.txt")   # FileNotFoundError is suppressed

print("Still running after the error.")

# Output:
# Suppressed: [Errno 2] No such file or directory: 'does_not_exist.txt'
# Still running after the error.

This is exactly what contextlib.suppress does internally.

Writing a Context Manager with @contextlib.contextmanager

Writing a class every time is overkill. The contextlib.contextmanager decorator lets you write a context manager as a simple generator function:

  • Everything before yield is the setup (__enter__).
  • The yield value is what gets bound to the as variable.
  • Everything after yield is the teardown (__exit__).
from contextlib import contextmanager

@contextmanager
def managed_file(path, mode="r"):
    print(f"Opening {path!r}")
    f = open(path, mode)
    try:
        yield f           # give the caller the file object
    finally:
        print(f"Closing {path!r}")
        f.close()         # always runs


with managed_file("data.txt", "w") as f:
    f.write("Written via contextmanager!\n")

# Opening 'data.txt'
# Closing 'data.txt'

The try/finally inside the generator ensures cleanup happens even if the body of the with block raises an exception. You must always wrap the yield in try/finally when there's something to clean up.

Practical Context Manager Patterns

Timer

from contextlib import contextmanager
import time

@contextmanager
def timer(label=""):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        tag = f"[{label}] " if label else ""
        print(f"{tag}Elapsed: {elapsed:.4f}s")


with timer("Sorting"):
    data = sorted(range(1_000_000), reverse=True)

# [Sorting] Elapsed: 0.0821s

Temporary directory

from contextlib import contextmanager
import tempfile
import shutil
import os

@contextmanager
def temp_directory():
    """Create a temp directory, yield its path, then delete it."""
    path = tempfile.mkdtemp()
    try:
        yield path
    finally:
        shutil.rmtree(path, ignore_errors=True)


with temp_directory() as tmpdir:
    test_file = os.path.join(tmpdir, "test.txt")
    with open(test_file, "w") as f:
        f.write("Temporary content")
    print(f"Temp dir: {tmpdir}")
    print(f"File exists: {os.path.exists(test_file)}")

print(f"Dir still exists after: {os.path.exists(tmpdir)}")
# False — deleted automatically

Changing directory temporarily

from contextlib import contextmanager
import os

@contextmanager
def working_directory(path):
    """Temporarily change the current working directory."""
    original = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(original)


print(os.getcwd())                    # C:\python-book
with working_directory("C:\\Users"):
    print(os.getcwd())                # C:\Users
print(os.getcwd())                    # C:\python-book (restored)

Redirecting stdout

from contextlib import contextmanager, redirect_stdout
import io

@contextmanager
def capture_output():
    """Capture everything printed inside the block."""
    buffer = io.StringIO()
    with redirect_stdout(buffer):
        yield buffer
    # buffer.getvalue() has everything that was printed


with capture_output() as output:
    print("Hello")
    print("World")
    for i in range(3):
        print(f"Line {i}")

captured = output.getvalue()
print(f"Captured {len(captured.splitlines())} lines:")
print(repr(captured))

Output:

Captured 5 lines:
'Hello\nWorld\nLine 0\nLine 1\nLine 2\n'

Database transaction

from contextlib import contextmanager
import sqlite3

@contextmanager
def transaction(connection):
    """Commit on success, rollback on failure."""
    try:
        yield connection
        connection.commit()
        print("Transaction committed.")
    except Exception as e:
        connection.rollback()
        print(f"Transaction rolled back: {e}")
        raise


conn = sqlite3.connect(":memory:")
conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT)")

# Successful transaction
with transaction(conn):
    conn.execute("INSERT INTO users (name) VALUES (?)", ("Alice",))
    conn.execute("INSERT INTO users (name) VALUES (?)", ("Bob",))

# Failed transaction — both inserts rolled back
try:
    with transaction(conn):
        conn.execute("INSERT INTO users (name) VALUES (?)", ("Carlos",))
        raise ValueError("Simulated error mid-transaction")
except ValueError:
    pass

rows = conn.execute("SELECT * FROM users").fetchall()
print(rows)   # [(1, 'Alice'), (2, 'Bob')] — Carlos never committed
conn.close()

Lock management

from contextlib import contextmanager
import threading

@contextmanager
def locked(lock):
    """Acquire a lock, yield, then release it."""
    lock.acquire()
    try:
        yield
    finally:
        lock.release()


# (Same as using: with lock:)
my_lock = threading.Lock()
shared_data = []

def append_safely(item):
    with locked(my_lock):
        shared_data.append(item)

contextlib Utilities You Should Know

Python's contextlib module has several ready-made tools:

from contextlib import (
    suppress,
    nullcontext,
    ExitStack,
    asynccontextmanager,
)

# suppress — ignore specific exceptions
from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("does_not_exist.txt")   # silently ignored
# equivalent to our SuppressErrors class from earlier


# nullcontext — a no-op context manager (useful for optional context managers)
def process(data, lock=None):
    ctx = lock if lock is not None else nullcontext()
    with ctx:
        # do something with data
        return data * 2

process(5)           # no lock
process(5, my_lock)  # with lock


# ExitStack — dynamically manage multiple context managers
files = ["a.txt", "b.txt", "c.txt"]

with ExitStack() as stack:
    handles = [stack.enter_context(open(f, "w")) for f in files]
    for i, handle in enumerate(handles):
        handle.write(f"Content of file {i}\n")
# All three files closed automatically — even if one write fails

ExitStack is particularly useful when you don't know at write time how many context managers you need.

Nested and Multiple Context Managers

You can use multiple context managers in one with statement:

# Old way — nested
with open("input.txt") as infile:
    with open("output.txt", "w") as outfile:
        outfile.write(infile.read())

# New way — comma-separated (Python 3.1+)
with open("input.txt") as infile, open("output.txt", "w") as outfile:
    outfile.write(infile.read())

Both are equivalent. The comma form is preferred — it's flatter and just as clear.

Project: A Robust Database Connection Manager

Let's build a complete context manager that handles SQLite connections with automatic transaction management, connection pooling awareness, and detailed logging.

"""
db.py — A context manager for robust SQLite access.
"""
from contextlib import contextmanager
import sqlite3
import logging
import time

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
log = logging.getLogger("db")


class Database:
    """
    Context manager wrapping a SQLite connection.

    Usage:
        with Database("myapp.db") as db:
            db.execute("INSERT INTO ...")
            rows = db.query("SELECT ...")
    """

    def __init__(self, path, timeout=30, isolation_level="DEFERRED"):
        self.path            = path
        self.timeout         = timeout
        self.isolation_level = isolation_level
        self._conn           = None

    # ── Context manager protocol ──────────────────────────────────────────────

    def __enter__(self):
        log.info(f"Connecting to {self.path!r}")
        self._conn = sqlite3.connect(
            self.path,
            timeout=self.timeout,
            isolation_level=self.isolation_level,
            check_same_thread=False,
        )
        self._conn.row_factory = sqlite3.Row   # rows behave like dicts
        self._conn.execute("PRAGMA journal_mode=WAL")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._conn:
            if exc_type is None:
                log.info("Committing transaction")
                self._conn.commit()
            else:
                log.warning(f"Rolling back due to: {exc_val}")
                self._conn.rollback()
            self._conn.close()
            self._conn = None
            log.info(f"Connection to {self.path!r} closed")
        return False   # don't suppress exceptions

    # ── Query helpers ─────────────────────────────────────────────────────────

    def execute(self, sql, params=()):
        """Run a write statement. Returns the cursor."""
        if self._conn is None:
            raise RuntimeError("Not inside a 'with' block.")
        return self._conn.execute(sql, params)

    def executemany(self, sql, seq_of_params):
        if self._conn is None:
            raise RuntimeError("Not inside a 'with' block.")
        return self._conn.executemany(sql, seq_of_params)

    def query(self, sql, params=()):
        """Run a SELECT and return a list of Row objects."""
        return self.execute(sql, params).fetchall()

    def query_one(self, sql, params=()):
        """Return the first matching row, or None."""
        return self.execute(sql, params).fetchone()

    def table_exists(self, name):
        row = self.query_one(
            "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?",
            (name,),
        )
        return row is not None

    # ── Schema helpers ────────────────────────────────────────────────────────

    @contextmanager
    def savepoint(self, name="sp"):
        """Nested transaction using SQLite savepoints."""
        self.execute(f"SAVEPOINT {name}")
        try:
            yield self
            self.execute(f"RELEASE SAVEPOINT {name}")
        except Exception:
            self.execute(f"ROLLBACK TO SAVEPOINT {name}")
            raise


# ── Demo ──────────────────────────────────────────────────────────────────────

def setup_schema(db):
    if not db.table_exists("users"):
        db.execute("""
            CREATE TABLE users (
                id    INTEGER PRIMARY KEY AUTOINCREMENT,
                name  TEXT    NOT NULL,
                email TEXT    UNIQUE NOT NULL,
                score REAL    DEFAULT 0
            )
        """)
        log.info("Created table 'users'")


def seed_data(db):
    users = [
        ("Alice",   "alice@example.com",   95.5),
        ("Bob",     "bob@example.com",     82.0),
        ("Carlos",  "carlos@example.com",  91.3),
        ("Diana",   "diana@example.com",   88.7),
    ]
    db.executemany(
        "INSERT OR IGNORE INTO users (name, email, score) VALUES (?, ?, ?)",
        users,
    )
    log.info(f"Seeded {len(users)} users")


if __name__ == "__main__":
    DB_PATH = ":memory:"

    # Normal operation — commits automatically
    with Database(DB_PATH) as db:
        setup_schema(db)
        seed_data(db)

        rows = db.query("SELECT name, score FROM users ORDER BY score DESC")
        print("\nTop scorers:")
        for row in rows:
            print(f"  {row['name']:<10} {row['score']}")

    # Savepoint — nested transaction
    with Database(DB_PATH) as db:
        setup_schema(db)   # table already exists, no-op

        try:
            with db.savepoint("add_user"):
                db.execute(
                    "INSERT INTO users (name, email, score) VALUES (?, ?, ?)",
                    ("Eve", "eve@example.com", 99.9),
                )
                # Simulate a conflict
                db.execute(
                    "INSERT INTO users (name, email, score) VALUES (?, ?, ?)",
                    ("Eve", "eve@example.com", 50.0),   # duplicate email
                )
        except sqlite3.IntegrityError as e:
            log.warning(f"Savepoint rolled back: {e}")

        # Eve was not added (savepoint rolled back)
        count = db.query_one("SELECT COUNT(*) AS n FROM users")
        print(f"\nUser count after conflict: {count['n']}")   # 4

What You Learned in This Chapter

  • Context managers guarantee setup and teardown code runs, even when exceptions occur.
  • The with statement calls __enter__() at the start and __exit__() at the end.
  • If __exit__ returns True, exceptions are suppressed; False lets them propagate.
  • Write context managers as classes (implement __enter__ and __exit__) or as generator functions with @contextmanager (wrap yield in try/finally).
  • Common patterns: timers, temp directories, working directory changes, stdout capture, database transactions, lock management.
  • contextlib.suppress — ignore specific exception types cleanly.
  • contextlib.nullcontext — a no-op for optional context managers.
  • contextlib.ExitStack — manage a dynamic number of context managers.
  • Multiple context managers can share one with line using commas.

What's Next?

Chapter 25 covers Functional Programming — map, filter, reduce, pure functions, immutability, and using functools to write code that is easier to reason about, test, and compose. You already know the building blocks from earlier chapters; this chapter ties them together into a coherent style.

© 2026 Abhilash Sahoo. Python: Zero to Hero.