Topic 14 of 64
with statement
Overview
Context managers ensure resources (files, database connections, locks) are properly acquired and released, even if an exception occurs. The with statement is Python's elegant solution for deterministic resource cleanup.
Syntax
python
# Using built-in context managers
with open("file.txt", "r") as f:
data = f.read() # file auto-closes after block
# Database connection
with db.get_connection() as conn:
cursor = conn.cursor()
cursor.execute("SELECT * FROM users")
# Multiple context managers
with open("input.txt") as inp, open("output.txt", "w") as out:
out.write(inp.read())
# Creating custom context managers
# Method 1: Class-based (__enter__ / __exit__)
class Timer:
def __enter__(self):
self.start = time.perf_counter()
return self
def __exit__(self, *args):
self.elapsed = time.perf_counter() - self.start
with Timer() as t:
expensive_operation()
print(f"Took {t.elapsed:.4f}s")
# Method 2: contextlib.contextmanager
from contextlib import contextmanager
@contextmanager
def managed_session(db):
session = db.create_session()
try:
yield session
session.commit()
except Exception:
session.rollback()
raise
finally:
session.close()Common Pitfalls
- The __exit__ method receives exception info — return True to suppress the exception, False (or None) to propagate it.
- contextlib.suppress() is a neat built-in: with suppress(FileNotFoundError): os.remove(file) — silently ignores specific exceptions.
- Interview tip: Context managers are implemented using Python's protocol system (__enter__/__exit__) — the with statement is syntactic sugar for calling these.
Real-World Example
Database transaction context manager for a banking system:
example
python
from contextlib import contextmanager
from typing import Generator
import psycopg2
@contextmanager
def transaction(conn) -> Generator:
"""Context manager that commits on success, rolls back on failure."""
cursor = conn.cursor()
try:
yield cursor
conn.commit()
except Exception as e:
conn.rollback()
raise
finally:
cursor.close()
def transfer_funds(from_id: int, to_id: int, amount: float) -> bool:
with get_db_connection() as conn:
with transaction(conn) as cur:
# Both operations must succeed or both fail
cur.execute(
"UPDATE accounts SET balance = balance - %s WHERE id = %s",
(amount, from_id)
)
cur.execute(
"UPDATE accounts SET balance = balance + %s WHERE id = %s",
(amount, to_id)
)
return True