Context Managers

0/5 in this phase0/54 across the roadmap

📖 Concept

Context managers ensure that resources are properly acquired and released, regardless of whether an operation succeeds or fails. The with statement guarantees cleanup code runs, even if exceptions occur — no more forgotten file.close() calls or unreleased locks.

The context manager protocol:

with expression as variable:
    # __enter__() is called, return value bound to variable
    body
# __exit__() is ALWAYS called (even on exception)

Two ways to create context managers:

Approach	Best for	Complexity
Class-based (`__enter__`/`__exit__`)	Complex state, reusable managers	More boilerplate
`@contextmanager` (generator-based)	Simple setup/teardown, quick one-offs	Less code

__exit__ parameters: When an exception occurs inside the with block, __exit__ receives the exception type, value, and traceback. Returning True from __exit__ suppresses the exception; returning False (or None) lets it propagate.

Nested context managers can be combined with contextlib.ExitStack for dynamic or variable-length resource management:

with ExitStack() as stack:
    files = [stack.enter_context(open(f)) for f in file_list]

Async context managers use __aenter__/__aexit__ and are entered with async with. The @asynccontextmanager decorator from contextlib simplifies creation. They're essential for managing async resources like database connections, HTTP sessions, and network sockets in asyncio code.

💻 Code Example

codeTap to expand ⛶

1# ============================================================
2# Class-based context manager
3# ============================================================
4import time
5import contextlib
6import logging
7import os
8import tempfile
9from typing import Optional
10
11logger = logging.getLogger(__name__)
12
13
14class Timer:
15    """Context manager that measures execution time."""
16
17    def __init__(self, label: str = "Block"):
18        self.label = label
19        self.elapsed: Optional[float] = None
20
21    def __enter__(self):
22        self._start = time.perf_counter()
23        return self  # Bound to 'as' variable
24
25    def __exit__(self, exc_type, exc_val, exc_tb):
26        self.elapsed = time.perf_counter() - self._start
27        logger.info(f"{self.label} took {self.elapsed:.4f}s")
28        return False  # Don't suppress exceptions
29
30
31with Timer("Data processing") as t:
32    total = sum(range(1_000_000))
33print(f"Elapsed: {t.elapsed:.4f}s")
34
35
36# ============================================================
37# Context manager with exception handling
38# ============================================================
39class DatabaseTransaction:
40    """Manage database transactions with automatic commit/rollback."""
41
42    def __init__(self, connection):
43        self.connection = connection
44
45    def __enter__(self):
46        self.connection.begin()
47        return self.connection
48
49    def __exit__(self, exc_type, exc_val, exc_tb):
50        if exc_type is not None:
51            # Exception occurred — rollback
52            self.connection.rollback()
53            logger.error(f"Transaction rolled back: {exc_val}")
54            return False  # Re-raise the exception
55        else:
56            # No exception — commit
57            self.connection.commit()
58            return False
59
60
61# Usage:
62# with DatabaseTransaction(conn) as db:
63#     db.execute("INSERT INTO users (name) VALUES (?)", ("Alice",))
64#     db.execute("UPDATE accounts SET balance = balance - 100 WHERE id = 1")
65# Automatically commits on success, rolls back on exception
66
67
68# ============================================================
69# @contextmanager — generator-based (simpler)
70# ============================================================
71@contextlib.contextmanager
72def temporary_directory(prefix: str = "tmp_"):
73    """Create a temp directory, yield it, then clean up."""
74    path = tempfile.mkdtemp(prefix=prefix)
75    logger.info(f"Created temp dir: {path}")
76    try:
77        yield path  # Everything before yield = __enter__
78    finally:
79        # Everything after yield = __exit__ (always runs)
80        import shutil
81        shutil.rmtree(path, ignore_errors=True)
82        logger.info(f"Cleaned up temp dir: {path}")
83
84
85with temporary_directory("myapp_") as tmpdir:
86    # Work with temporary directory
87    filepath = os.path.join(tmpdir, "data.txt")
88    with open(filepath, "w") as f:
89        f.write("temporary data")
90# Directory is automatically deleted here
91
92
93# ============================================================
94# @contextmanager with exception handling
95# ============================================================
96@contextlib.contextmanager
97def managed_resource(name: str):
98    """Acquire and release a named resource."""
99    print(f"Acquiring {name}")
100    resource = {"name": name, "active": True}
101    try:
102        yield resource
103    except Exception as e:
104        print(f"Error in {name}: {e}")
105        resource["error"] = str(e)
106        raise  # Re-raise after logging
107    finally:
108        resource["active"] = False
109        print(f"Released {name}")
110
111
112# ============================================================
113# ExitStack for dynamic context management
114# ============================================================
115def process_multiple_files(file_paths: list[str]):
116    """Open and process a variable number of files."""
117    with contextlib.ExitStack() as stack:
118        # Dynamically enter multiple context managers
119        files = [
120            stack.enter_context(open(path, "r"))
121            for path in file_paths
122        ]
123        # All files are open here; ALL will be closed on exit
124        for f in files:
125            print(f.readline())
126
127
128# ============================================================
129# Suppressing specific exceptions
130# ============================================================
131# Instead of try/except/pass:
132with contextlib.suppress(FileNotFoundError):
133    os.remove("nonexistent_file.txt")
134    # No error even though file doesn't exist
135
136
137# ============================================================
138# Reentrant and reusable context managers
139# ============================================================
140@contextlib.contextmanager
141def indent_logger(level: int = 1):
142    """Context manager that indents log output."""
143    prefix = "  " * level
144    original_factory = logging.getLogRecordFactory()
145
146    def indented_factory(*args, **kwargs):
147        record = original_factory(*args, **kwargs)
148        record.msg = f"{prefix}{record.msg}"
149        return record
150
151    logging.setLogRecordFactory(indented_factory)
152    try:
153        yield
154    finally:
155        logging.setLogRecordFactory(original_factory)
156
157
158# ============================================================
159# Async context manager
160# ============================================================
161import asyncio
162
163
164class AsyncDatabasePool:
165    """Async context manager for database connection pools."""
166
167    def __init__(self, dsn: str, min_size: int = 2, max_size: int = 10):
168        self.dsn = dsn
169        self.min_size = min_size
170        self.max_size = max_size
171        self.pool = None
172
173    async def __aenter__(self):
174        # Simulate async pool creation
175        print(f"Creating pool for {self.dsn}")
176        await asyncio.sleep(0.1)  # Simulate connection time
177        self.pool = {"dsn": self.dsn, "connections": self.min_size}
178        return self.pool
179
180    async def __aexit__(self, exc_type, exc_val, exc_tb):
181        print(f"Closing pool for {self.dsn}")
182        await asyncio.sleep(0.05)  # Simulate graceful shutdown
183        self.pool = None
184        return False
185
186
187@contextlib.asynccontextmanager
188async def async_http_session(base_url: str):
189    """Async context manager for HTTP session lifecycle."""
190    session = {"base_url": base_url, "active": True}
191    print(f"Opening session to {base_url}")
192    try:
193        yield session
194    finally:
195        session["active"] = False
196        print(f"Closed session to {base_url}")
197
198
199# Usage:
200# async def main():
201#     async with AsyncDatabasePool("postgresql://...") as pool:
202#         async with async_http_session("https://api.example.com") as http:
203#             # Both resources managed
204#             pass

🏋️ Practice Exercise

Exercises:

Write a class-based context manager FileBackup that copies a file to a .bak before entering the block, and restores it from backup if an exception occurs during the block. Delete the backup on successful exit.
Implement @contextmanager-based change_directory(path) that changes the working directory on entry and restores the original directory on exit, even if an exception occurs.
Build an ExitStack-based function that opens N database connections, N file handles, and N network sockets (simulated), processes data from all of them, and ensures all resources are released on exit.
Create a context manager timeout(seconds) using signal.alarm (Unix) that raises TimeoutError if the block takes too long. Handle the cleanup properly in __exit__.
Write an async context manager rate_limiter(max_concurrent) using asyncio.Semaphore that limits how many coroutines can execute the body concurrently. Test it with 20 simulated API calls limited to 5 concurrent.

⚠️ Common Mistakes

Forgetting to return False (or not returning at all) from __exit__ and accidentally suppressing exceptions. Only return True if you intentionally want to swallow the exception. Returning None (implicit) correctly propagates exceptions.
Using @contextmanager but forgetting the try/finally around yield. Without finally, the cleanup code after yield won't run if an exception occurs inside the with block — defeating the entire purpose of the context manager.
Yielding more than once in a @contextmanager generator. The generator must yield exactly once — the yield separates setup from teardown. Multiple yields cause RuntimeError.
Not understanding that __exit__ is called even when __enter__ assigns to a variable that's never used. The as clause is optional; the context manager's lifecycle (enter/exit) runs regardless.
Creating context managers that hold resources indefinitely. If a context manager acquires a database connection in __enter__, holding the with block open for a long time (e.g., waiting for user input) starves the connection pool. Keep with blocks short.

💼 Interview Questions

🎤 Mock Interview

Practice a live interview for Context Managers

Was this topic helpful?

← PreviousGenerators & Iterators Next →Metaclasses