🔵 Intermediate

Context Managers

📖 Lesson 23 ⏱ 35 min 🧪 5 questions 💻 3 exercises

🎯 Learning Objectives

Understand what a context manager is and why it exists
Use the with statement correctly
Build context managers with __enter__ / __exit__
Build context managers with @contextmanager and yield
Use contextlib utilities: suppress, redirect_stdout, ExitStack
Apply context managers to real-world patterns: timers, transactions, temporary files

The Problem They Solve

Many operations require a matching pair of setup and teardown actions — open/close, lock/unlock, start/stop, connect/disconnect. The risk is that an exception in the middle leaves the teardown unexecuted:

# ❌ Fragile — close() may never be called if an exception occurs
f = open("data.txt")
data = process(f)   # what if this raises?
f.close()           # skipped on exception!

# ✓ Context manager guarantees close() — even on exception
with open("data.txt") as f:
    data = process(f)
# f.close() called here, always

problem.py

A context manager is any object that implements the context manager protocol: __enter__ (setup) and __exit__ (teardown). The with statement calls these automatically and guarantees teardown even when exceptions occur.

The `with` Statement in Depth

# Full form — 'as target' binds the return value of __enter__
with expression as target:
    body

# What Python actually does:
manager = expression
target  = manager.__enter__()
try:
    body
except:
    if not manager.__exit__(*sys.exc_info()):
        raise      # re-raise if __exit__ returns falsy
else:
    manager.__exit__(None, None, None)

with_expansion.py

Multiple contexts in one line

# Open two files simultaneously — both are closed even if one raises
with open("input.txt", "r") as src, open("output.txt", "w") as dst:
    dst.write(src.read().upper())

# Python 3.10+ parenthesised form — easier to read with many managers
with (
    open("input.txt",  "r", encoding="utf-8") as src,
    open("output.txt", "w", encoding="utf-8") as dst,
):
    dst.write(src.read().upper())

multiple_with.py

Class-Based Context Managers

Implement __enter__ and __exit__ on any class to make it a context manager:

import time

class Timer:
    """Measure the wall-clock time of a block of code."""

    def __enter__(self):
        self.start   = time.perf_counter()
        self.elapsed = None
        return self          # bound to the 'as' variable

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start
        # exc_type/val/tb are None when no exception occurred
        return False         # don't suppress exceptions

with Timer() as t:
    result = sum(range(10_000_000))

print(f"Result: {result:,}")
print(f"Time:   {t.elapsed:.4f}s")

timer_class.py

import threading

class ManagedLock:
    """Acquire a threading lock, release it on exit."""

    def __init__(self, lock):
        self._lock = lock

    def __enter__(self):
        self._lock.acquire()
        return self._lock

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._lock.release()
        return False

lock = threading.Lock()

with ManagedLock(lock):
    # Only one thread can be here at a time
    print("Protected section")

managed_lock.py

__exit__ receives three arguments: exc_type, exc_val, exc_tb — all None if no exception occurred. Return True to suppress the exception; return False (or None) to propagate it.

`@contextmanager` — Generator-Based

contextlib.contextmanager lets you write a context manager as a simple generator function — often cleaner than a full class:

from contextlib import contextmanager

@contextmanager
def timer():
    """Measure elapsed time of a with-block."""
    import time
    start = time.perf_counter()
    try:
        yield          # execution of the with-block happens here
    finally:
        elapsed = time.perf_counter() - start
        print(f"Elapsed: {elapsed:.4f}s")

with timer():
    total = sum(range(10_000_000))

contextmanager_basic.py

from contextlib import contextmanager

@contextmanager
def managed_file(path, mode="r", **kwargs):
    """Open a file, yield it, always close it."""
    f = open(path, mode, **kwargs)
    try:
        yield f
    finally:
        f.close()

with managed_file("notes.txt", "w", encoding="utf-8") as f:
    f.write("Hello from a generator-based context manager!\n")

managed_file.py

How it works: Everything before yield is the __enter__ logic. The value after yield is what gets bound to the as variable. Everything after yield (typically in a finally block) is the __exit__ logic.

Yielding a value

from contextlib import contextmanager

@contextmanager
def temp_directory():
    """Create a temporary directory; delete it on exit."""
    import tempfile, shutil, os
    tmpdir = tempfile.mkdtemp()
    try:
        yield tmpdir          # bound to 'as' variable
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)

with temp_directory() as tmpdir:
    filepath = f"{tmpdir}/test.txt"
    with open(filepath, "w") as f:
        f.write("temporary data")
    print(f"Working in: {tmpdir}")
# Directory is deleted here, even if an exception occurred

temp_dir.py

Handling exceptions inside @contextmanager

from contextlib import contextmanager

@contextmanager
def database_transaction(db):
    """Commit on success, rollback on exception."""
    db.begin()
    try:
        yield db
    except Exception:
        db.rollback()
        raise          # re-raise after rollback
    else:
        db.commit()    # only if no exception

# Usage:
# with database_transaction(db) as conn:
#     conn.execute("INSERT INTO ...")

transaction.py

Useful `contextlib` Utilities

suppress — silence specific exceptions

from contextlib import suppress
from pathlib import Path

# Equivalent to try/except/pass — but clearer intent
with suppress(FileNotFoundError):
    Path("maybe_missing.txt").unlink()

# Multiple exception types
with suppress(KeyError, AttributeError):
    value = data["missing_key"].strip()

suppress.py

redirect_stdout / redirect_stderr

import io
from contextlib import redirect_stdout

# Capture print() output as a string
buffer = io.StringIO()
with redirect_stdout(buffer):
    print("Hello")
    print("World")

output = buffer.getvalue()
print(repr(output))   # 'Hello\nWorld\n'

# Useful for testing functions that print to stdout
def noisy_function():
    print("lots of debug output")

with redirect_stdout(io.StringIO()):
    noisy_function()   # output discarded

redirect.py

nullcontext — optional context manager

from contextlib import nullcontext

def process(data, lock=None):
    """Use a lock if provided; otherwise no-op context manager."""
    ctx = lock if lock is not None else nullcontext()
    with ctx:
        return [x * 2 for x in data]

# No lock — nullcontext is a no-op
result = process([1, 2, 3])
print(result)   # [2, 4, 6]

# With a lock — uses it transparently
import threading
result = process([1, 2, 3], lock=threading.Lock())
print(result)   # [2, 4, 6]

nullcontext.py

ExitStack — dynamic context managers

from contextlib import ExitStack

# Open a variable number of files — can't use a fixed with statement
filenames = ["a.txt", "b.txt", "c.txt"]

with ExitStack() as stack:
    files = [
        stack.enter_context(open(f, "w", encoding="utf-8"))
        for f in filenames
    ]
    for i, fh in enumerate(files):
        fh.write(f"File {i}\n")
# All files are closed here, even if one raised an exception

# Also useful for conditional context managers
def process(path, verbose=False):
    with ExitStack() as stack:
        f = stack.enter_context(open(path, encoding="utf-8"))
        if verbose:
            stack.enter_context(Timer())   # add timer only when verbose
        return f.read()

exit_stack.py

ExitStack is the go-to tool when the number of context managers is not known until runtime. It also works as a manual cleanup stack: call stack.callback(func) to register any zero-argument cleanup function.

Real-World Patterns

Atomic file write

from contextlib import contextmanager
from pathlib import Path
import os

@contextmanager
def atomic_write(path, mode="w", **kwargs):
    """Write to a temp file; replace target only on clean exit."""
    tmp = Path(str(path) + ".tmp")
    try:
        with open(tmp, mode, **kwargs) as f:
            yield f
        tmp.replace(path)   # atomic rename — only reached if no exception
    except:
        tmp.unlink(missing_ok=True)
        raise

with atomic_write("config.json", encoding="utf-8") as f:
    import json
    json.dump({"version": 2, "debug": False}, f, indent=2)
# config.json is updated atomically

atomic_write.py

Temporary environment variable

from contextlib import contextmanager
import os

@contextmanager
def env_var(name, value):
    """Temporarily set an environment variable, restore original on exit."""
    original = os.environ.get(name)
    os.environ[name] = str(value)
    try:
        yield
    finally:
        if original is None:
            del os.environ[name]
        else:
            os.environ[name] = original

with env_var("DEBUG", "1"):
    print(os.environ["DEBUG"])   # 1

print(os.environ.get("DEBUG"))   # None (or original value)

env_var.py

Indented printing

from contextlib import contextmanager

_indent = 0

@contextmanager
def indent(level=2):
    global _indent
    _indent += level
    try:
        yield
    finally:
        _indent -= level

def iprint(msg):
    print(" " * _indent + msg)

iprint("Start")
with indent():
    iprint("Level 1")
    with indent():
        iprint("Level 2")
    iprint("Back to 1")
iprint("End")
# Start
#   Level 1
#     Level 2
#   Back to 1
# End

indent_print.py

Class vs `@contextmanager`

	Class (`__enter__/__exit__`)	`@contextmanager`
Verbosity	More code	Concise
State between enter/exit	Easy — store on `self`	Local variables work fine
Reusable as base class	Yes — inheritable	No
Exception handling	`__exit__` receives all info	Use `try/except/finally` around `yield`
Best for	Complex managers, reusable libraries	Simple, one-off managers

Best Practices

Always use with for resources that need cleanup (files, locks, connections, temp directories).
Keep with blocks focused. Only the code that needs the resource should be inside the block.
Use @contextmanager for simple cases; use a class when you need inheritance or complex state.
Always wrap the yield in a try/finally inside @contextmanager — otherwise exceptions skip the cleanup.
Only return True from __exit__ if you deliberately want to suppress the exception — be explicit about this.
Use contextlib.suppress instead of try/except/pass for conciseness.

from contextlib import contextmanager

# ❌ Missing finally — exception leaks without cleanup
@contextmanager
def bad_cm(resource):
    resource.open()
    yield resource
    resource.close()   # skipped if an exception occurs in the with-block!

# ✓ Always use try/finally
@contextmanager
def good_cm(resource):
    resource.open()
    try:
        yield resource
    finally:
        resource.close()   # guaranteed

best_practice.py

Primary sources: Python Docs — The with statement · Python Docs — contextlib · PEP 343 — The "with" Statement

🤖

Ask your AI tutor! Not sure whether to use a class or @contextmanager for your use case? Want to see how ExitStack handles cleanup for a variable number of resources? Great patterns to work through with concrete code.

💻 Exercises

01 Retry Context Manager

Write a @contextmanager called retry(times, exceptions) that re-runs the body of the with block up to times attempts whenever one of the specified exceptions is raised. After all attempts are exhausted, re-raise the last exception.

import random

for attempt in retry(times=5, exceptions=(ValueError,)):
    with attempt:
        if random.random() < 0.7:
            raise ValueError("random failure")
        print("Success!")
        break

Hint: yield attempt objects that know whether to continue looping.

Show solution

from contextlib import contextmanager

class _Attempt:
    """Sentinel object yielded to the with-block."""
    def __enter__(self): return self
    def __exit__(self, *args): return False

def retry(times, exceptions=(Exception,)):
    last_exc = None
    for attempt in range(times):
        try:
            yield _Attempt()
            return   # success — stop iterating
        except exceptions as e:
            last_exc = e
            print(f"Attempt {attempt + 1}/{times} failed: {e}")
    raise last_exc

import random, types

# Make it a generator-based loop
def retry_loop(times, exceptions=(Exception,)):
    last_exc = None
    for i in range(times):
        try:
            yield i
            return
        except exceptions as e:
            last_exc = e
            print(f"Attempt {i+1}/{times} failed: {e}")
    if last_exc:
        raise last_exc

# Simpler contextmanager version
@contextmanager
def retry_ctx(times=3, exceptions=(Exception,), delay=0):
    import time
    last_exc = None
    for attempt in range(1, times + 1):
        try:
            yield attempt
            return   # clean exit — done
        except exceptions as e:
            last_exc = e
            print(f"Attempt {attempt}/{times} failed: {e}")
            if attempt < times and delay:
                time.sleep(delay)
    raise last_exc

# Usage
random.seed(42)
for _ in range(10):
    try:
        with retry_ctx(times=5, exceptions=(ValueError,)) as attempt:
            if random.random() < 0.7:
                raise ValueError("random failure")
            print(f"Succeeded on attempt {attempt}")
        break
    except ValueError:
        print("All attempts failed")

02 Managed Database Connection

Build a DatabaseConnection class-based context manager that simulates a database connection with transaction support:

__enter__: "connect" (print a message) and return self
execute(sql): append SQL to an internal log; raise RuntimeError if not connected
__exit__: if no exception → commit (print log); if exception → rollback; always disconnect
A queries property returning the list of executed SQL strings

Show solution

class DatabaseConnection:
    def __init__(self, url):
        self.url       = url
        self._log      = []
        self._connected = False

    @property
    def queries(self):
        return list(self._log)

    def execute(self, sql):
        if not self._connected:
            raise RuntimeError("Not connected")
        self._log.append(sql)
        print(f"  EXEC: {sql}")

    def __enter__(self):
        self._connected = True
        self._log.clear()
        print(f"Connected to {self.url}")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            print(f"COMMIT — {len(self._log)} statement(s)")
        else:
            print(f"ROLLBACK — {exc_val}")
            self._log.clear()
        self._connected = False
        print("Disconnected")
        return False   # don't suppress exceptions

# Happy path
with DatabaseConnection("postgres://localhost/app") as db:
    db.execute("INSERT INTO users VALUES (1, 'Alice')")
    db.execute("UPDATE stats SET logins = logins + 1")

print(db.queries)

# Error path — rollback
try:
    with DatabaseConnection("postgres://localhost/app") as db:
        db.execute("DELETE FROM users WHERE id = 99")
        raise ValueError("Something went wrong")
except ValueError:
    pass  # already rolled back
print(db.queries)   # []

03 Temporary Config Override

Write a @contextmanager called config_override(config_dict, **overrides) that temporarily patches a config dictionary with new values and restores all original values on exit — even if an exception occurs. Demonstrate it with a simple app config dict.

config = {"debug": False, "log_level": "WARNING", "port": 8080}

with config_override(config, debug=True, log_level="DEBUG"):
    print(config)   # debug=True, log_level='DEBUG', port=8080

print(config)       # fully restored

Show solution

from contextlib import contextmanager

@contextmanager
def config_override(cfg, **overrides):
    """Temporarily patch a dict; restore originals on exit."""
    # Save originals (or sentinel for keys that didn't exist)
    _MISSING = object()
    saved = {k: cfg.get(k, _MISSING) for k in overrides}

    # Apply overrides
    cfg.update(overrides)
    try:
        yield cfg
    finally:
        # Restore
        for k, original in saved.items():
            if original is _MISSING:
                cfg.pop(k, None)
            else:
                cfg[k] = original

config = {"debug": False, "log_level": "WARNING", "port": 8080}

print("Before:", config)

with config_override(config, debug=True, log_level="DEBUG", new_key="added"):
    print("Inside:", config)
    # {"debug": True, "log_level": "DEBUG", "port": 8080, "new_key": "added"}

print("After: ", config)
# {"debug": False, "log_level": "WARNING", "port": 8080}
# 'new_key' is removed because it didn't exist before

# Works even when an exception occurs
try:
    with config_override(config, debug=True):
        raise RuntimeError("crash!")
except RuntimeError:
    pass

print("After exception:", config["debug"])   # False — restored

Context Managers

🎯 Learning Objectives

The Problem They Solve

The with Statement in Depth

Multiple contexts in one line

Class-Based Context Managers

@contextmanager — Generator-Based

Yielding a value

Handling exceptions inside @contextmanager

Useful contextlib Utilities

suppress — silence specific exceptions

redirect_stdout / redirect_stderr

nullcontext — optional context manager

ExitStack — dynamic context managers

Real-World Patterns

Atomic file write

Temporary environment variable

Indented printing

Class vs @contextmanager

Best Practices

💻 Exercises

The `with` Statement in Depth

`@contextmanager` — Generator-Based

Useful `contextlib` Utilities

Class vs `@contextmanager`