On this article, you’ll find out how Python context managers work beneath the floor, and methods to apply them to actual cleanup issues nicely past file dealing with.
Subjects we’ll cowl embody:
- How the
withassertion’s protocol works, together with what__enter__()and__exit__()truly do and why the return worth of__exit__()issues. - The best way to construct customized context managers utilizing a category, the
@contextmanagerdecorator, and customary library instruments likeTemporaryDirectoryandpatch.dict(). - The best way to handle a runtime-variable variety of sources 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 take a look at 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 correct 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 document 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 cause they’ll see, and three GB of extracted information 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 foul document 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, be taught that the file will get closed, and cease there. Information are the least attention-grabbing factor a context supervisor can handle. What the with assertion actually supplies 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 go away it exterior the block? Every of the 4 issues the analysis script units up is a short lived contract. Every could be 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 sources that unwinds in the correct 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 title after as, the block runs, after which __exit__() is 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 title receives no matter __enter__() returns, which is commonly the supervisor itself however doesn’t need 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 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 top 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 info mannequin.
Construct a Timer That Can’t Neglect 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 decisions 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 to be like 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 option to the traceback, and twelve seconds is the quantity you wish to see when the actual query is whether or not the batch remodel was hanging moderately than crashing.
perf_counter() is the correct clock right here. It has the very best decision out there for brief durations, and on CPython it can’t go backwards. Its absolute worth means nothing; solely the distinction between two calls does, which the category hides by design. Maintain 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 truly prices the twenty minutes. A Timer line isn’t a benchmark. The second somebody begins evaluating two Timer strains throughout completely different machines, the subsequent step is timeit or a profiler moderately than extra decimals.
Use @contextmanager for One Setup-Yield-Cleanup Stream
A category is the correct 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: every little thing earlier than yield runs on entry, every little thing 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)     strive:         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 preserve 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 moderately than on the road after the yield. A strive/besides that catches the exception to log it has to reraise it, or the supervisor has quietly turn into 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 nicely 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 every little thing 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 move as immediately’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. It is usually simpler to examine whereas the block is working: print tmp and open it in one other shell. When a run wants its artifacts stored for a autopsy, Python 3.12 added a delete=False parameter for precisely that, which beats commenting out an rmtree.
One platform notice 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, out there since Python 3.10, downgrades that to a best-effort sweep, which is okay in a script and a foul 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 thus far 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 strive/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 sources in the correct order is the issue ExitStack exists for, and the documentation describes it as a option 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 strains in zip(*shards):             out.writelines(strains) |
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 don’t have any supervisor in any respect, corresponding to a tool launch or a notification. Plain callbacks can’t 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 every little thing 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() is known as, not when the stack object is rubbish collected, so an ExitStack that escapes its with block is a leak. And an internal callback that suppresses an exception adjustments what the outer ones see, the identical as nested with blocks would.
One trustworthy caveat: if the one cause a chunk 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 Adjustments Reversible
The setting variable within the opening is a state change moderately than a useful resource, and state adjustments are the place cleanup bugs are most costly, as a result of nothing crashes. The cache path is just 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 bought, 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, “machine”: “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 moderately 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, momentary paths, take a look at credentials, and the small configuration override that may in any other case be a worldwide one perform units and one other forgets. The title comes from the testing module, however nothing about it’s test-only. It is usually a lot tougher to get unsuitable than a hand-written outdated = d.get(okay) 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 youngster 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 without delay. Atmosphere variables and the present working listing are the 2 sorts of momentary state that behave badly below concurrency for precisely that cause. The usual library’s personal contextlib.chdir() carries a warning to that impact: it adjustments international state, so it’s not appropriate for many threaded or async code, or for turbines that yield whereas it’s energetic. Reverse international state when the method is doing one factor at a time. When it isn’t, move the trail or the setting as an argument as an alternative.
Maintain 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 sources 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 until the block is supposed to proceed after a failure. contextlib.suppress() is reserved for the particular errors the place persevering with is thought to be proper.
| Sample | Finest match | Keep away from when | Main 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 until suppression is the intent |
@contextmanager quiet(logger) |
One linear setup, yield, cleanup circulate; cleanup in lastly | State have to be learn after the block, or cleanup has branches | Precisely one yield; a logged exception have to be reraised |
built-in supervisor TemporaryDirectory patch.dict |
The usual library already owns the lifecycle | The lifecycle differs from the library’s (preserve artifacts, partial restore) | World state (environ, cwd) is process-wide and unsafe below threads or async |
ExitStack enter_context() callback() |
The variety of sources 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, momentary artifacts, machine handles, periods, 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 to be like 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 best way out. Cleanup that lives on the backside of the file is a promise the joyful path retains and the error path breaks. The error path is the one which runs at two within the morning.
