8.2. Consider contextlib and with Statements for Reusable try/finally Behavior¶
The with statement in Python is used to indicate when code is running in a special context. For example, mutual-exclusion locks (see Item 54: “Use Lock to Prevent Data Races in Threads”) can be used in with statements to indicate that the indented code block runs only while the lock is held:
Click here to view code image
>>> from threading import Lock
>>>
>>> lock = Lock()
>>>
>>> with lock:
>>> # Do something while maintaining an invariant
>>> ...
The example above is equivalent to this try/finally construction because the Lock class properly enables the with statement (see Item 65: “Take Advantage of Each Block in try/except/else/finally” for more about try/finally):
Click here to view code image
lock.acquire() try:
# Do something while maintaining an invariant ...
- finally:
lock.release()
The with statement version of this is better because it eliminates the need to write the repetitive code of the try/finally construction, and it ensures that you don’t forget to have a corresponding release call for every acquire call.
It’s easy to make your objects and functions work in with statements by using the contextlib built-in module. This module contains the contextmanager decorator (see Item 26: “Define Function Decorators with functools.wraps” for background), which lets a simple function be used in with statements. This is much easier than defining a new class with the special methods enter and exit (the standard way).
For example, say that I want a region of code to have more debug logging sometimes. Here, I define a function that does logging at two severity levels:
Click here to view code image
>>> import logging
>>>
>>> def my_function():
>>> logging.debug('Some debug data')
>>> logging.error('Error log here')
>>> logging.debug('More debug data')
The default log level for my program is WARNING, so only the error message will print to screen when I run the function:
my_function()
>>>
Error log here
I can elevate the log level of this function temporarily by defining a context manager. This helper function boosts the logging severity level before running the code in the with block and reduces the logging severity level afterward:
Click here to view code image
>>> from contextlib import contextmanager
>>>
>>> @contextmanager
>>> def debug_logging(level):
>>> logger = logging.getLogger()
>>> old_level = logger.getEffectiveLevel()
>>> logger.setLevel(level)
>>> try:
>>> yield
>>> finally:
>>> logger.setLevel(old_level)
The yield expression is the point at which the with block’s contents will execute (see Item 30: “Consider Generators Instead of Returning Lists” for background). Any exceptions that happen in the with block will be re-raised by the yield expression for you to catch in the helper function (see Item 35: “Avoid Causing State Transitions in Generators with throw” for how that works).
Now, I can call the same logging function again but in the debug_logging context. This time, all of the debug messages are printed to the screen during the with block. The same function running outside the with block won’t print debug messages:
- with debug_logging(logging.DEBUG):
print('* Inside:') my_function()
print('* After:') my_function()
>>>
* Inside:
Some debug data
Error log here
More debug data
* After:
Error log here
Using with Targets
The context manager passed to a with statement may also return an object. This object is assigned to a local variable in the as part of the compound statement. This gives the code running in the with block the ability to directly interact with its context.
For example, say I want to write a file and ensure that it’s always closed correctly. I can do this by passing open to the with statement. open returns a file handle for the as target of with, and it closes the handle when the with block exits:
Click here to view code image
>>> with open('my_output.txt', 'w') as handle:
>>> handle.write('This is some data!')
This approach is more Pythonic than manually opening and closing the file handle every time. It gives you confidence that the file is eventually closed when execution leaves the with statement. By highlighting the critical section, it also encourages you to reduce the amount of code that executes while the file handle is open, which is good practice in general.
To enable your own functions to supply values for as targets, all you need to do is yield a value from your context manager. For example, here I define a context manager to fetch a Logger instance, set its level, and then yield it as the target:
Click here to view code image
>>> @contextmanager
>>> def log_level(level, name):
>>> logger = logging.getLogger(name)
>>> old_level = logger.getEffectiveLevel()
>>> logger.setLevel(level)
>>> try:
>>> yield logger
>>> finally:
>>> logger.setLevel(old_level)
Calling logging methods like debug on the as target produces output because the logging severity level is set low enough in the with block on that specific Logger instance. Using the logging module directly won’t print anything because the default logging severity level for the default program logger is WARNING:
Click here to view code image
>>> with log_level(logging.DEBUG, 'my-log') as logger:
>>> logger.debug(f'This is a message for {logger.name}!')
>>> logging.debug('This will not print')
DEBUG:my-log:This is a message for my-log!
After the with statement exits, calling debug logging methods on the Logger named ‘my-log’ will not print anything because the default logging severity level has been restored. Error log messages will always print:
Click here to view code image
>>> logger = logging.getLogger('my-log')
>>> logger.debug('Debug will not print')
>>> logger.error('Error will print')
ERROR:my-log:Error will print
Later, I can change the name of the logger I want to use by simply updating the with statement. This will point the Logger that’s the as target in the with block to a different instance, but I won’t have to update any of my other code to match:
Click here to view code image
>>> with log_level(logging.DEBUG, 'other-log') as logger:
>>> logger.debug(f'This is a message for {logger.name}!')
>>> logging.debug('This will not print')
DEBUG:other-log:This is a message for other-log!
This isolation of state and decoupling between creating a context and acting within that context is another benefit of the with statement.
8.2.1. Things to Remember¶
✦ The with statement allows you to reuse logic from try/finally blocks and reduce visual noise.
✦ The contextlib built-in module provides a contextmanager decorator that makes it easy to use your own functions in with statements.
✦ The value yielded by context managers is supplied to the as part of the with statement. It’s useful for letting your code directly access the cause of a special context.