Context Managers Beyond Opening Files in Python

In this article, you will learn how Python context managers work beneath the surface, and how to apply them to real cleanup problems well beyond file handling. Topics we will cover include: How the with statement’s protocol works, including what __enter__() and __exit__() actually do and why the return value of __exit__() matters. How to build custom context managers using a class, the @contextmanager decorator, and standard library tools like TemporaryDirectory and patch.dict(). How to manage a runtime-variable number of resources safely using ExitStack, and when to wrap it in a class of its own. The evaluation script does a fair amount of housekeeping before it evaluates anything. There is a scratch directory for the extracted test set, an environment variable that points the model cache at a faster disk, a library logger turned down to WARNING because it otherwise floods the console, and a stopwatch. The bottom of the file undoes all of it in the right order: the directory is deleted, the variable restored, the logger reset, the elapsed time printed. The cleanup at the bottom runs every time the script succeeds. The day a malformed record makes evaluate() raise halfway through, none of it runs. Whoever uses the machine next gets a stale cache path, a logger that has gone quiet for no reason they can see, and 3 GB of extracted files nobody remembers creating. scratch = Path("scratch"); scratch.mkdir(exist_ok=True)os.environ["MODEL_CACHE"] = "/fast/cache"logging.getLogger("transformers").setLevel(logging.WARNING)start = time.perf_counter()scores = evaluate(model, extract(archive, scratch)) # raises on a bad recordprint(f"took {time.perf_counter() - start:.1f}s") # none of this runsshutil.rmtree(scratch)del os.environ["MODEL_CACHE"]logging.getLogger("transformers").setLevel(logging.INFO) Most Python developers meet context managers through with open(path) as f, learn that the file gets closed, and stop there. Files are the least interesting thing a context manager can manage. What the with statement really provides is a block with a guaranteed exit. Setup happens at the top, cleanup at the bottom, and the cleanup runs whether the middle succeeded or threw. If cleanup has to happen even when the block fails, why leave it outside the block? Each of the four things the evaluation script sets up is a temporary contract. Each can be a context manager: a timer that cannot forget to stop, a directory that removes itself, a state change that reverses itself, and a stack of resources that unwinds in the right order. That is the progression here, from the protocol underneath with through the standard library’s shortcuts to a rule for choosing between them. The With Statement Is a Boundary, Not a File Trick The language reference describes the with statement as a sequence, and the sequence is worth holding in your head. Every pattern below is a variation on it. Python evaluates the expression after with to get a context manager and calls its __enter__(). Whatever __enter__() returned is bound to the name after as, the block runs, and then __exit__() is called. The last step is the whole point. It happens on a normal exit, on a return from inside the block, and on an exception, so no path through the block can skip the cleanup written in __exit__(). Two details trip people up, and both concern return values. First, as name receives whatever __enter__() returns, which is often the manager itself but doesn’t have to be. open() returns the file, and TemporaryDirectory() returns a path string. Second, __exit__() receives three arguments describing the exception, or three None values on a clean exit, and its own return value decides what happens next. The data model is precise about this. Return a true value and the exception is suppressed. Return False, or fall off the end and return None, and the exception propagates after cleanup. class Announce: def __enter__(self): print("enter") return "the value bound after as" def __exit__(self, exc_type, exc_value, traceback): print("exit", exc_type) # None on a clean exit return False # let any exception propagatewith Announce() as x: print(x) The suppression rule is the one to get right first. Getting it wrong is silent. A __exit__() that ends with return True turns every exception inside the block into nothing, and so does one that returns the result of a cleanup call which happens to be truthy. The script keeps running, the scores are missing, and the traceback that would have explained why was eaten at the block boundary. The custom managers below never return anything truthy from __exit__(). Whatever went wrong inside the block stays the caller’s problem, which is where it belongs. The guarantee in one picture: both routes out of the block go through __exit__(), and on the exception route its return value is the only thing standing between the error and the caller. Credit: original diagram based on the Python language reference for the with statement and the data model. Build a Timer That Cannot Forget to Stop The stopwatch in the opening script is the smallest case of the pattern. Two calls to time.perf_counter() bracket the work, and the second call is exactly the kind of line an exception skips. A class with __enter__() and __exit__() moves that second call somewhere it cannot be skipped. 123456789101112131415161718192021 import timeclass Timer: def __init__(self, label): self.label = label self.elapsed = None def __enter__(self): self._start = time.perf_counter() return self def __exit__(self, exc_type, exc_value, traceback): self.elapsed = time.perf_counter() - self._start status = "failed" if exc_type else "ok" print(f"{self.label}: {self.elapsed:.3f}s ({status})") return False # never swallow the block's exceptionwith Timer("feature extraction") as t: features = extract(batch)print(t.elapsed) # still readable after the block Two choices in that class do the work. __enter__() returns the timer itself, which is why t.elapsed can be read once the block has finished, and __exit__() looks at exc_type only to choose between ok and failed for the label. __exit__() never keeps the exception. So a feature-extraction run that dies still prints feature extraction: 12.418s (failed) on its way to the traceback, and twelve seconds is the number you want to see when the real question is whether the batch transform was hanging rather than crashing. perf_counter() is the right clock here. It has the highest resolution available for short durations, and on CPython it cannot go backwards. Its absolute value means nothing; only the difference between two calls does, which the class hides by design. Keep the ambition of the class in proportion, though. It measures one block of wall-clock time in one process and prints a line. That is instrumentation, and it is useful for finding out which stage of a preprocessing job actually costs the twenty minutes. A Timer line is not a benchmark. The moment someone starts comparing two Timer lines across different machines, the next step is timeit or a profiler rather than more decimals. Use @contextmanager for One Setup-Yield-Cleanup Flow A class is the right shape when the manager has state a caller wants to read back, as the timer does. When the whole job is to set one thing, run the block, and put the thing back, a class is ceremony, and @contextmanager is the shortcut. It takes a generator function and wraps it in the same __enter__()/__exit__() guarantees, with the generator’s shape doing the work: everything before yield runs on entry, everything after it runs on exit. import loggingfrom contextlib import contextmanager@contextmanagerdef quiet(logger_name, level=logging.WARNING): logger = logging.getLogger(logger_name) previous = logger.level logger.setLevel(level) try: yield logger finally: logger.setLevel(previous)with quiet("transformers"): scores = evaluate(model, eval_set) The chatty-logger fix from the opening now has a shape that survives failure. The previous level is captured before the change, the change is made, the block runs at the yield, and the finally puts the level back. Two rules from the documentation keep this working. The generator must yield exactly one value, because that yield is the block, and a second one raises RuntimeError on exit. And if the block raises, the exception is reraised inside the generator at the yield. That is why the cleanup sits in finally rather than on the line after the yield. A try/except that catches the exception to log it has to reraise it, or the manager has quietly become a suppressor. The decorator gives one more thing away for free. @contextmanager builds on ContextDecorator, so quiet(...) can also sit above a function definition as a decorator, with the whole function body as the block. Machine Learning Mastery’s guide to production Python decorators covers where that reads well and where it hides too much. Let Temporary Directories Clean Their Own Mess The scratch directory in the opening is the case the standard library already owns, and the fix is to stop owning it yourself. tempfile.TemporaryDirectory() creates a uniquely named directory and hands the path to the as target. When the block ends, on either path out, the directory and everything inside it are removed. import tempfilefrom pathlib import Path# before: a fixed path, and a cleanup line that only runs on successscratch = Path("scratch")scratch.mkdir(exist_ok=True)scores = evaluate(extract(archive, scratch))shutil.rmtree(scratch)# after: the directory's lifetime is the blockwith tempfile.TemporaryDirectory(prefix="eval-") as tmp: scores = evaluate(extract(archive, Path(tmp)))# tmp is gone here, whether evaluate() raised or not The before version has three problems beyond the skipped cleanup. Two runs on the same machine share scratch/. A leftover file from yesterday’s archive can pass as today’s data. A crash leaves the debris in the project tree, where it gets committed. The after version puts the artifacts in a path that is unique per run and outside the repository. It is also easier to inspect while the block is running: print tmp and open it in another shell. When a run needs its artifacts kept for a post-mortem, Python 3.12 added a delete=False parameter for exactly that, which beats commenting out an rmtree. One platform note earns a sentence. On Windows, a file still open inside the directory makes removal fail with PermissionError at exit. Passing ignore_cleanup_errors=True, available since Python 3.10, downgrades that to a best-effort sweep, which is fine in a script and a bad habit in a library. The proper fix is to close what you opened before the outer block ends, usually by nesting a with for the file inside the one for the directory. Combine Dynamic Resources With ExitStack Every example so far manages one thing. The evaluation script’s harder relative is the merge step that opens one file per shard, and the number of shards is whatever the previous job produced. Nesting with statements needs the count at write time. A loop that calls open() and appends to a list needs a try/finally that closes each handle in reverse, and it must not lose an exception raised by the third close() while the second is already running. Closing a runtime-sized set of resources in the right order is the problem ExitStack exists for, and the documentation describes it as a way to programmatically combine other context managers and cleanup functions, especially those driven by input data. from contextlib import ExitStackdef merge_shards(paths, out_path): with ExitStack() as stack: shards = [stack.enter_context(open(p, encoding="utf-8")) for p in paths] out = stack.enter_context(open(out_path, "w", encoding="utf-8")) stack.callback(print, f"merged {len(paths)} shards into {out_path}") for lines in zip(*shards): out.writelines(lines) enter_context() calls the manager’s __enter__() and registers its __exit__() on the stack. It works for any context manager, including the Timer and TemporaryDirectory above. callback() registers a plain function for cases that have no manager at all, such as a device release or a notification. Plain callbacks cannot suppress exceptions, because they never see the exception details. At the end of the block, the stack unwinds backwards through everything it registered, so the output file closes first and the shards follow, last-opened first — exactly the order a hand-nested set of with statements would have produced. Two details matter in production. The unwinding happens when the block ends or close() is called, not when the stack object is garbage collected, so an ExitStack that escapes its with block is a leak. And an inner callback that suppresses an exception changes what the outer ones see, the same as nested with blocks would. One honest caveat: if the only reason a piece of code works is that ExitStack is holding a dozen handles, locks, and callbacks together, the stack is compensating for a missing abstraction. A ShardSet class with its own __enter__() and __exit__(), built on ExitStack internally, usually reads better than the loop. The caller then sees one resource with one lifetime instead of a list of things to remember. Make State Changes Reversible The environment variable in the opening is a state change rather than a resource, and state changes are where cleanup bugs are most expensive, because nothing crashes. The cache path is simply wrong for the next job, silently. For anything dictionary-shaped, the standard library has the reversal written already, though it hides in unittest.mock as patch.dict(). The documentation’s description is the whole feature: patch a dictionary, then restore it to its original state after the block. Any mapping will do, as long as items can be got, set, and deleted. import osfrom unittest.mock import patchwith patch.dict(os.environ, {"MODEL_CACHE": "/fast/cache", "TOKENIZERS_PARALLELISM": "false"}): scores = evaluate(model, eval_set)# MODEL_CACHE is back to its old value; TOKENIZERS_PARALLELISM is unset again if it was unset beforeCONFIG = {"batch_size": 64, "device": "cuda"}with patch.dict(CONFIG, {"batch_size": 8}): smoke_test() # runs with batch_size 8, and CONFIG is restored afterwards The restore is a real restore rather than a delete. A key that existed goes back to its old value, and a key that did not exist is removed. That covers feature flags, temporary paths, test credentials, and the small configuration override that would otherwise be a global one function sets and another forgets. The name comes from the testing module, but nothing about it is test-only. It is also much harder to get wrong than a hand-written old = d.get(k) block that has to handle the missing-key case itself. The scope is what to be careful about. os.environ is a mapping over the current process’s environment, captured when os is imported and pushed back through putenv() on every change. A patch.dict() on it changes what this process, and any child it spawns during the block, will see. Nothing already running elsewhere notices. Within the process, though, the change is visible to every thread and every coroutine at once. Environment variables and the current working directory are the two kinds of temporary state that behave badly under concurrency for exactly that reason. The standard library’s own contextlib.chdir() carries a warning to that effect: it changes global state, so it is not suitable for most threaded or async code, or for generators that yield while it is active. Reverse global state when the process is doing one thing at a time. When it isn’t, pass the path or the setting as an argument instead. Keep Cleanup Visible Where It Matters The selection rule is short. Write a class when the manager has state a caller needs afterwards, like the timer’s elapsed. Write a @contextmanager generator when the job is one setup, one block, and one cleanup, like the logger level. Use the standard library’s manager when it already owns the lifecycle, as TemporaryDirectory does for scratch artifacts and patch.dict does for reversible mappings. Reach for ExitStack when the set of resources is only known at runtime, and wrap it in a class of its own once it holds more than a handful. Whatever the shape, __exit__() returns False unless the block is meant to continue after a failure. contextlib.suppress() is reserved for the specific errors where continuing is known to be right. Pattern Best fit Avoid when Primary caveat class + __enter__/__exit__ Timer The manager has state a caller reads back afterwards, or several methods of its own One setup, one block, one cleanup and nothing else __exit__() must return False or None unless suppression is the intent @contextmanager quiet(logger) One linear setup, yield, cleanup flow; cleanup in finally State must be read after the block, or cleanup has branches Exactly one yield; a logged exception must be reraised built-in manager TemporaryDirectory patch.dict The standard library already owns the lifecycle The lifecycle differs from the library’s (keep artifacts, partial restore) Global state (environ, cwd) is process-wide and unsafe under threads or async ExitStack enter_context() callback() The number of resources is only known at runtime It is hiding a resource object that should exist Unwinds at block end or close(), never on garbage collection The four shapes and when each one is the right call. Credit: original table based on the Python contextlib, tempfile, time and unittest.mock documentation. The hidden payoff is the one the opening script was missing. Profiling scopes, temporary artifacts, device handles, sessions, and environment overrides all have a lifetime. The with block makes that lifetime local and visible instead of spreading it between line 12 and line 140. Machine Learning Mastery’s survey of Python concepts every AI engineer should know lists context managers as one item, and this is what the item looks like once it is doing real work. A reader opening the module should be able to see, from one block, what gets acquired, what can fail inside it, and what gets restored on the way out. Cleanup that lives at the bottom of the file is a promise the happy path keeps and the error path breaks. The error path is the one that runs at two in the morning. No comments yet.