On this article, you’ll learn the way Python context managers work beneath the floor, and methods to apply them to actual cleanup issues properly past file dealing with.
Matters we are going to cowl embody:
- How the
withassertion’s protocol works, together with what__enter__()and__exit__()really do and why the return worth of__exit__()issues. - How you can construct customized context managers utilizing a category, the
@contextmanagerdecorator, and commonplace library instruments likeTemporaryDirectoryandpatch.dict(). - How you can handle a runtime-variable variety of assets safely utilizing
ExitStack, and when to wrap it in a category of its personal.
The analysis script does a good quantity of housekeeping earlier than it evaluates something. There’s a scratch listing for the extracted check set, an setting variable that factors the mannequin cache at a sooner disk, a library logger turned all the way down to WARNING as a result of it in any other case floods the console, and a stopwatch. The underside of the file undoes all of it in the proper order: the listing is deleted, the variable restored, the logger reset, the elapsed time printed. The cleanup on the backside runs each time the script succeeds. The day a malformed file makes consider() elevate midway by, none of it runs. Whoever makes use of the machine subsequent will get a stale cache path, a logger that has gone quiet for no motive they’ll see, and three GB of extracted recordsdata no person remembers creating.
|
scratch = Path(“scratch”); scratch.mkdir(exist_ok=True) os.environ[“MODEL_CACHE”] = “/quick/cache” logging.getLogger(“transformers”).setLevel(logging.WARNING)  begin = time.perf_counter()  scores = consider(mannequin, extract(archive, scratch))  # raises on a nasty file print(f“took {time.perf_counter() – begin:.1f}s”)    # none of this runs  shutil.rmtree(scratch) del os.environ[“MODEL_CACHE”] logging.getLogger(“transformers”).setLevel(logging.INFO) |
Most Python builders meet context managers by with open(path) as f, study that the file will get closed, and cease there. Recordsdata are the least attention-grabbing factor a context supervisor can handle. What the with assertion actually offers is a block with a assured exit. Setup occurs on the high, cleanup on the backside, and the cleanup runs whether or not the center succeeded or threw. If cleanup has to occur even when the block fails, why depart it outdoors the block? Every of the 4 issues the analysis script units up is a brief contract. Every generally is a context supervisor: a timer that can’t overlook to cease, a listing that removes itself, a state change that reverses itself, and a stack of assets that unwinds in the proper order. That’s the development right here, from the protocol beneath with by the usual library’s shortcuts to a rule for selecting between them.
The With Assertion Is a Boundary, Not a File Trick
The language reference describes the with assertion as a sequence, and the sequence is value holding in your head. Each sample under is a variation on it. Python evaluates the expression after with to get a context supervisor and calls its __enter__(). No matter __enter__() returned is sure to the identify after as, the block runs, after which __exit__() known as. The final step is the entire level. It occurs on a standard exit, on a return from contained in the block, and on an exception, so no path by the block can skip the cleanup written in __exit__().
Two particulars journey folks up, and each concern return values. First, as identify receives no matter __enter__() returns, which is commonly the supervisor itself however doesn’t must 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 clear exit, and its personal return worth decides what occurs subsequent. The knowledge mannequin is exact about this. Return a real worth and the exception is suppressed. Return False, or fall off the tip and return None, and the exception propagates after cleanup.
|
class Announce:     def __enter__(self):         print(“enter”)         return “the worth sure after as”      def __exit__(self, exc_type, exc_value, traceback):         print(“exit”, exc_type)        # None on a clear exit         return False                  # let any exception propagate  with Announce() as x:     print(x) |
The suppression rule is the one to get proper first. Getting it unsuitable is silent. A __exit__() that ends with return True turns each exception contained in the block into nothing, and so does one which returns the results of a cleanup name which occurs to be truthy. The script retains working, the scores are lacking, and the traceback that may have defined why was eaten on the block boundary. The customized managers under by no means return something truthy from __exit__(). No matter went unsuitable contained in the block stays the caller’s drawback, which is the place it belongs.
The assure in a single image: each routes out of the block undergo __exit__(), and on the exception route its return worth is the one factor standing between the error and the caller. Credit score: authentic diagram based mostly on the Python language reference for the with assertion and the information mannequin.
Construct a Timer That Can’t Overlook to Cease
The stopwatch within the opening script is the smallest case of the sample. Two calls to time.perf_counter() bracket the work, and the second name is precisely the form of line an exception skips. A category with __enter__() and __exit__() strikes that second name someplace it can’t be skipped.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
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         standing = “failed” if exc_type else “okay”         print(f“{self.label}: {self.elapsed:.3f}s ({standing})”)         return False                  # by no means swallow the block’s exception  with Timer(“function extraction”) as t:     options = extract(batch)  print(t.elapsed)                      # nonetheless readable after the block |
Two selections in that class do the work. __enter__() returns the timer itself, which is why t.elapsed will be learn as soon as the block has completed, and __exit__() appears at exc_type solely to decide on between okay and failed for the label. __exit__() by no means retains the exception. So a feature-extraction run that dies nonetheless prints function extraction: 12.418s (failed) on its solution to the traceback, and twelve seconds is the quantity you need to see when the actual query is whether or not the batch remodel was hanging fairly than crashing.
perf_counter() is the proper clock right here. It has the best decision obtainable for brief durations, and on CPython it can not go backwards. Its absolute worth means nothing; solely the distinction between two calls does, which the category hides by design. Hold the ambition of the category in proportion, although. It measures one block of wall-clock time in a single course of and prints a line. That’s instrumentation, and it’s helpful for locating out which stage of a preprocessing job really prices the twenty minutes. A Timer line isn’t a benchmark. The second somebody begins evaluating two Timer traces throughout completely different machines, the subsequent step is timeit or a profiler fairly than extra decimals.
Use @contextmanager for One Setup-Yield-Cleanup Circulate
A category is the proper form when the supervisor has state a caller desires to learn again, because the timer does. When the entire job is to set one factor, run the block, and put the factor again, a category is ceremony, and @contextmanager is the shortcut. It takes a generator perform and wraps it in the identical __enter__()/__exit__() ensures, with the generator’s form doing the work: all the pieces earlier than yield runs on entry, all the pieces after it runs on exit.
|
import logging from contextlib import contextmanager  @contextmanager def quiet(logger_name, degree=logging.WARNING):     logger = logging.getLogger(logger_name)     earlier = logger.degree     logger.setLevel(degree)     attempt:         yield logger     lastly:         logger.setLevel(earlier)  with quiet(“transformers”):     scores = consider(mannequin, eval_set) |
The chatty-logger repair from the opening now has a form that survives failure. The earlier degree is captured earlier than the change, the change is made, the block runs on the yield, and the lastly places the extent again. Two guidelines from the documentation hold this working. The generator should yield precisely one worth, as a result of that yield is the block, and a second one raises RuntimeError on exit. And if the block raises, the exception is reraised contained in the generator on the yield. That’s the reason the cleanup sits in lastly fairly than on the road after the yield. A attempt/besides that catches the exception to log it has to reraise it, or the supervisor has quietly grow to be a suppressor.
The decorator provides yet one more factor away totally free. @contextmanager builds on ContextDecorator, so quiet(...) can even sit above a perform definition as a decorator, with the entire perform physique because the block. Machine Studying Mastery’s information to manufacturing Python decorators covers the place that reads properly and the place it hides an excessive amount of.
Let Non permanent Directories Clear Their Personal Mess
The scratch listing within the opening is the case the usual library already owns, and the repair is to cease proudly owning it your self. tempfile.TemporaryDirectory() creates a uniquely named listing and fingers the trail to the as goal. When the block ends, on both path out, the listing and all the pieces inside it are eliminated.
|
import tempfile from pathlib import Path  # earlier than: a set path, and a cleanup line that solely runs on success scratch = Path(“scratch”) scratch.mkdir(exist_ok=True) scores = consider(extract(archive, scratch)) shutil.rmtree(scratch)  # after: the listing’s lifetime is the block with tempfile.TemporaryDirectory(prefix=“eval-“) as tmp:     scores = consider(extract(archive, Path(tmp)))  # tmp is gone right here, whether or not consider() raised or not |
The earlier than model has three issues past the skipped cleanup. Two runs on the identical machine share scratch/. A leftover file from yesterday’s archive can go as at the moment’s knowledge. A crash leaves the particles within the mission tree, the place it will get dedicated. The after model places the artifacts in a path that’s distinctive per run and out of doors the repository. Additionally it is simpler to examine whereas the block is working: print tmp and open it in one other shell. When a run wants its artifacts saved for a autopsy, Python 3.12 added a delete=False parameter for precisely that, which beats commenting out an rmtree.
One platform observe earns a sentence. On Home windows, a file nonetheless open contained in the listing makes removing fail with PermissionError at exit. Passing ignore_cleanup_errors=True, obtainable since Python 3.10, downgrades that to a best-effort sweep, which is ok in a script and a nasty behavior in a library. The right repair is to shut what you opened earlier than the outer block ends, often by nesting a with for the file contained in the one for the listing.
Mix Dynamic Assets With ExitStack
Each instance to this point manages one factor. The analysis script’s tougher relative is the merge step that opens one file per shard, and the variety of shards is regardless of the earlier job produced. Nesting with statements wants the depend at write time. A loop that calls open() and appends to a listing wants a attempt/lastly that closes every deal with in reverse, and it should not lose an exception raised by the third shut() whereas the second is already working. Closing a runtime-sized set of assets in the proper order is the issue ExitStack exists for, and the documentation describes it as a solution to programmatically mix different context managers and cleanup capabilities, particularly these pushed by enter knowledge.
|
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 traces in zip(*shards):             out.writelines(traces) |
enter_context() calls the supervisor’s __enter__() and registers its __exit__() on the stack. It really works for any context supervisor, together with the Timer and TemporaryDirectory above. callback() registers a plain perform for instances that haven’t any supervisor in any respect, equivalent to a tool launch or a notification. Plain callbacks can not suppress exceptions, as a result of they by no means see the exception particulars. On the finish of the block, the stack unwinds backwards by all the pieces it registered, so the output file closes first and the shards observe, last-opened first — precisely the order a hand-nested set of with statements would have produced. Two particulars matter in manufacturing. The unwinding occurs when the block ends or shut() known as, not when the stack object is rubbish collected, so an ExitStack that escapes its with block is a leak. And an interior callback that suppresses an exception adjustments what the outer ones see, the identical as nested with blocks would.
One trustworthy caveat: if the one motive a bit of code works is that ExitStack is holding a dozen handles, locks, and callbacks collectively, the stack is compensating for a lacking abstraction. A ShardSet class with its personal __enter__() and __exit__(), constructed on ExitStack internally, often reads higher than the loop. The caller then sees one useful resource with one lifetime as an alternative of a listing of issues to recollect.
Make State Modifications Reversible
The setting variable within the opening is a state change fairly than a useful resource, and state adjustments are the place cleanup bugs are costliest, as a result of nothing crashes. The cache path is solely unsuitable for the subsequent job, silently. For something dictionary-shaped, the usual library has the reversal written already, although it hides in unittest.mock as patch.dict(). The documentation’s description is the entire function: patch a dictionary, then restore it to its authentic state after the block. Any mapping will do, so long as objects will be received, set, and deleted.
|
import os from unittest.mock import patch  with patch.dict(os.environ, {“MODEL_CACHE”: “/quick/cache”, “TOKENIZERS_PARALLELISM”: “false”}):     scores = consider(mannequin, eval_set) # MODEL_CACHE is again to its outdated worth; TOKENIZERS_PARALLELISM is unset once more if it was unset earlier than  CONFIG = {“batch_size”: 64, “gadget”: “cuda”}  with patch.dict(CONFIG, {“batch_size”: 8}):     smoke_test()                      # runs with batch_size 8, and CONFIG is restored afterwards |
The restore is an actual restore fairly than a delete. A key that existed goes again to its outdated worth, and a key that didn’t exist is eliminated. That covers function flags, short-term paths, check credentials, and the small configuration override that may in any other case be a world one perform units and one other forgets. The identify comes from the testing module, however nothing about it’s test-only. Additionally it is a lot tougher to get unsuitable than a hand-written outdated = d.get(ok) block that has to deal with the missing-key case itself.
The scope is what to watch out about. os.environ is a mapping over the present course of’s setting, captured when os is imported and pushed again by putenv() on each change. A patch.dict() on it adjustments what this course of, and any baby it spawns through the block, will see. Nothing already working elsewhere notices. Inside the course of, although, the change is seen to each thread and each coroutine directly. Surroundings variables and the present working listing are the 2 sorts of short-term state that behave badly below concurrency for precisely that motive. The usual library’s personal contextlib.chdir() carries a warning to that impact: it adjustments international state, so it isn’t appropriate for many threaded or async code, or for mills that yield whereas it’s energetic. Reverse international state when the method is doing one factor at a time. When it isn’t, go the trail or the setting as an argument as an alternative.
Hold Cleanup Seen The place It Issues
The choice rule is brief. Write a category when the supervisor has state a caller wants afterwards, just like the timer’s elapsed. Write a @contextmanager generator when the job is one setup, one block, and one cleanup, just like the logger degree. Use the usual library’s supervisor when it already owns the lifecycle, as TemporaryDirectory does for scratch artifacts and patch.dict does for reversible mappings. Attain for ExitStack when the set of assets is just recognized at runtime, and wrap it in a category of its personal as soon as it holds greater than a handful. Regardless of the form, __exit__() returns False except the block is supposed to proceed after a failure. contextlib.suppress() is reserved for the precise errors the place persevering with is thought to be proper.
| Sample | Greatest match | Keep away from when | Major caveat |
|---|---|---|---|
class + __enter__/__exit__ Timer |
The supervisor has state a caller reads again afterwards, or a number of strategies of its personal | One setup, one block, one cleanup and nothing else | __exit__() should return False or None except suppression is the intent |
@contextmanager quiet(logger) |
One linear setup, yield, cleanup movement; cleanup in lastly | State should be learn after the block, or cleanup has branches | Precisely one yield; a logged exception should be reraised |
built-in supervisor TemporaryDirectory patch.dict |
The usual library already owns the lifecycle | The lifecycle differs from the library’s (hold artifacts, partial restore) | International state (environ, cwd) is process-wide and unsafe below threads or async |
ExitStack enter_context() callback() |
The variety of assets is just recognized at runtime | It’s hiding a useful resource object that ought to exist | Unwinds at block finish or shut(), by no means on rubbish assortment |
The hidden payoff is the one the opening script was lacking. Profiling scopes, short-term artifacts, gadget handles, classes, and setting overrides all have a lifetime. The with block makes that lifetime native and visual as an alternative of spreading it between line 12 and line 140. Machine Studying Mastery’s survey of Python ideas each AI engineer ought to know lists context managers as one merchandise, and that is what the merchandise appears like as soon as it’s doing actual work.
A reader opening the module ought to be capable to see, from one block, what will get acquired, what can fail inside it, and what will get restored on the way in which out. Cleanup that lives on the backside of the file is a promise the pleased path retains and the error path breaks. The error path is the one which runs at two within the morning.
