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 exceptionThe 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 ignoredcontextlib.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 automaticallyThe 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 exceptionSolution
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 NoneTests
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"