bgilbert / contextmanager.md
In Python, a context manager is an object that can be used in a with statement. Here’s a context manager that reports the total wall-clock time spent inside a with block:
import time class Timer(object): def __init__(self, msg): self._msg = msg def __enter__(self): # time.monotonic() requires Python >= 3.3 self._start = time.monotonic() def __exit__(self, exc_type, exc_value, exc_traceback): if exc_type: print('Failed: <>: <>'.format(self._msg, exc_value)) else: print('<>: <> s'.format(self._msg, time.monotonic() - self._start))
with Timer("doing stuff"): for i in range(1000000): pass
which produces this output:
doing stuff: 0.04281306266784668 s
Context managers can also handle exceptions from inside the block they guard. Timer does this. If we do:
with Timer("doing stuff"): raise ValueError('ack')
Failed: doing stuff: ack Traceback (most recent call last): File "test.py", line 20, in raise ValueError('ack') ValueError: ack
Because __exit__() doesn’t return True , the exception continues to propagate, so we see both the Failed message and the traceback. If __exit__() returned True , the exception would be swallowed by the context manager and we would see only the Failed line.
Context from the context manager
Let’s say we wanted the context manager to return some information, perhaps the granularity of the timer. We can modify __enter__() like this:
def __enter__(self): # time functions require Python >= 3.3 self._start = time.monotonic() return time.clock_getres(time.CLOCK_MONOTONIC)
with Timer("doing stuff") as resolution: print('Resolution: <>'.format(resolution)) for i in range(1000000): pass
Resolution: 1e-09 doing stuff: 0.043778783998277504 s
The contextmanager decorator
Writing a context manager involves some boilerplate: the __enter__() and __exit__() methods, and possibly an __init__() method to handle any arguments. There’s an easier way: the contextlib.contextmanager decorator. We can rewrite Timer like this:
import contextlib import time @contextlib.contextmanager def Timer(msg): start = time.monotonic() try: yield time.clock_getres(time.CLOCK_MONOTONIC) except BaseException as e: print('Failed: <>: <>'.format(msg, e)) raise else: print('<>: <> s'.format(msg, time.monotonic() - start))
Now Timer is just a generator that yields once (yielding the value to be bound by the with statement). The yield expression raises any exception thrown from the block, which the context manager can then handle or not.
One use for context managers is ensuring that an object (file, network connection, database handle) is closed when leaving a block. contextlib.closing() does this: its __exit__() method calls another object’s close() . Use it like this:
import contextlib with contextlib.closing(open("/etc/passwd")) as fh: lines = fh.readlines()
But that seems a bit verbose. A better approach: many objects implement __enter__() and __exit__() , so you can use them to manage themselves. The implementation of file (the type returned by open() ) has something like this:
def __enter__(self): return self def __exit__(self, exc_type, exc_value, exc_traceback): self.close()
with open("/etc/passwd") as fh: lines = fh.readlines()
Reimplement the contextmanager decorator. You’ll need to use some of the more obscure features of generators.
Python Context Manager Exception Handling and Retrying
One case use context manager to handle exceptions during execution of the with statement as can be seen in the snippet below. This is useful for example for rolling back database transactions in case of an exception, where the database connections can be retrieved from and returned to a connection pool.
from contextlib import contextmanager @contextmanager def managed_resource(param): try: print(f"db connection acquired param>") yield "this is the connection" except Exception as e: print(f"rolling back due to ex e>") raise finally: print("returning db connection to a pool") with managed_resource("this is param value") as connection: print(connection) raise ValueError("hello")
Execution of above prints below.
db connection acquired this is param value this is the connection rolling back due to ex hello returning db connection to a pool Traceback (most recent call last): File ". /test/test.py", line 36, in raise ValueError("hello") ValueError: hello
Without The With
What happens when we try to use managed_resource without with ?
Nothing. Returned object is only holds a reference to a generator and has __enter__() and __exit__() methods. The context object will only return a value from the generator upon __enter__() call, and will run the rest of the code after yield during __exit__() call.
Alternative Error Handling In Exit Method
Instead of the contextmanager wrapper you can implement __enter__ and __exit__ methods on our custom object like below:
class ManagedResource: def __init__(self, param): self.param = param def __enter__(self): print(f"db connection acquired self.param>") return self def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is not None: print(f"rolling back due to ex exc_val>") print("returning db connection to a pool") with ManagedResource("this is param value") as connection: print(connection) raise ValueError("hello")
Execution of above prints below.
db connection acquired this is param value rolling back due to ex hello returning db connection to a pool Traceback (most recent call last): File "/home/vackosar/src/vackosar.github.io/test/test.py", line 42, in raise ValueError("hello") ValueError: hello
Retries Using Context Manager Are Not Possible — Here Is An Alternative
Do you need to retry when for example storage is momentarily not available? However, You cannot implement retries of the code wrapped in with . You simply cannot execute yield multiple times. Instead, you can pass a callback to retry method containing a exponential backoff for loop below.
def retry_on_exception( fun: Callable, args=(), kwargs: Optional[dict] = None, retry_exceptions: Tuple = (Exception,), max_retries: int = 3, base_sleep_secs: float = 3.0, ): if kwargs is None: kwargs = dict() ex = None for retry in range(max_retries + 1): try: val = fun(*args, **kwargs) return val except retry_exceptions as e: ex = e sleep_secs = base_sleep_secs * 2 ** retry sleep(sleep_secs) raise RuntimeError(f"Too many retries (max_retries>) of fun.__name__>") from ex
Other Useful Posts
Did you know that you can implement functional ForEach in Bash?
Created on 08 Apr 2020. Updated on: 06 Jun 2022.