Context Managers Beyond Opening Files in Python

The stopwatch in the opening script is the smallest case of the pattern.

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 record
print(f"took {time.perf_counter() - start:.1f}s") # none of this runs
shutil.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 propagate

with 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.

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.

import time

class 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 exception

with 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 logging
from contextlib import contextmanager

@contextmanager
def 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 rera raised 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 tempfile
from pathlib import Path

# before: a fixed path, and a cleanup line that only runs on success
scratch = Path("scratch")
scratch.mkdir(exist_ok=True)
scores = evaluate(extract(archive, scratch))
shutil.rmtree(scratch)

# after: the directory's lifetime is the block
with 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 ExitStack

def 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 os
from unittest.mock import patch

with patch.dict(os.environ, {"MODEL_CACHE": "/fast/cache", "TOKENIZERS_PARALLELISM": "false"}):
    scores = evaluate(model, eval_set)