Software Engineer
Context Managers in Python: Using the "with" statement
LearnDataSci is reader-supported. When you purchase through links on our site, earned commissions help support our team of writers, researchers, and designers at no extra cost to you.
Context managers are used to set up and tear down temporary contexts, establish and resolve custom settings, and acquire and release resources. The open() function for opening files is one of the most familiar examples of a context manager.
Context managers sandwich code blocks between two distinct pieces of logic:
- The enter logic - this runs right before the nested code block executes
- The exit logic - this runs right after the nested code block is done.
The most common way you'll work with context managers is by using the with statement.
The with statement
A with statement is the primary method used to call a context manager. The general syntax is as follows:
with SomeContextManager as context_variable:
# do stuff with context_variable
This code works in the following order:
SomeContextManagerexecutes its enter (setup) logic before the indented code runs.SomeContextManagerbinds a value tocontext_variable, which can be used in the indented code- Inner code block runs:
# do stuff. SomeContextManagerexecutes its exit (cleanup) logic.
So why is this useful? In the next section, we'll discuss why this programming pattern comes in handy and why it's worth making your context managers from time to time.
Why use a context manager
Context managers keep our codebases much cleaner because they encapsulate administrative boilerplate and separate it from the business logic.
Additionally, context managers are structured to carry out their exit methods regardless of what happens in the code block they frame. So even if something goes wrong in the managed block, the context manager ensures the deallocations are performed and the default settings are restored.
Let's give a solid example. Think about operating on a file without using with, like in the following block.
f = None
try:
f = open('random.txt', 'r')
contents = f.read()
# do stuff with content
except Exception as e:
print(e)
finally:
if f:
f.close()The first thing to note is that we must always close an open file. The finally block would perform the close even if an error occurred. If we had to do this try-except-finally logic every time we wanted to work with a file we'd have a lot of duplicate code.
Luckily, Python's built-in open() is a context manager. Therefore, using a with statement, we can program the same logic like this:
with open('random.txt', 'r') as f:
contents = f.read()
# do stuff with contentsHere, open()'s enter method opens the file and returns a file object. The as keyword binds the returned value to f, and we use f to read the contents of random.txt. At the end of the execution of the inner code block, the exit method runs and closes the file.
We can check whether f is actually closed (with does not define a variable scope, we can access the variables it created from outside the statement).
print(f.closed)It's evident from this simple example that context managers allow us to make our code cleaner and more reusable.
Python defines several other context managers in the standard library, but it also allows programmers to define context managers of their own.
In the next section, we will work on defining custom context managers. We will first work on the simple function-based implementation and later move on to the slightly more complicated class-based definitions.
Creating Context Managers
Function-Based Implementation
The standard library provides contextlib, a module containing with statement utilities. One important helper is the @contextmanager decorator, which lets us define function-based context managers.
Let's try it out by making a context manager function that swaps the case of a string:
from contextlib import contextmanager
@contextmanager
def example_cm(string_input):
print('Setup logic\n')
swapped = string_input.swapcase()
try:
yield swapped
except ValueError as e:
print('An error occurred...')
finally:
print('\nTeardown logic\n')
print('End of context manager\n')
with example_cm('the MAJESTIC squirrel') as swapped_string:
# Managed code
print(swapped_string)The @contextmanager decorator expects a yield statement in the function body. yield divides the function into three parts:
- The expressions above
yieldare executed right before the code that is managed. - The managed code runs at
yield. Whatever is yielded makes up the context variable—swapped_stringin our case. - The expressions below
yieldare executed after the managed code is run.
Notice that we have placed yield within a try-except-finally block. This is not enforced, yet, it is good practice. This way, if an error occurs in the managed code the context manager will carry out its exit logic no matter what.
Look what happens if we raise a ValueError inside the with statement:
with example_cm('RaNdOm') as swapped_string:
raise(ValueError)The except block within example_cm handled the exception, and the finally block ensured we saw the exit message.
Notice that the "End of context manager" string is still printed. We're specifically catching a ValueError, but let's see what happens when an unhandled error is raised:
with example_cm('RaNdOm') as swapped_string:
raise(NameError)In this case, the finally block still executed regardless of the uncaught error, but the "End of context manager" printout didn't make it. This exemplifies why a finally block is helpful in many situations.
For more flexibility in creating context managers, we'll now introduce the class-based implementations.
Class-based implementation
Python defines a context management protocol that dictates that any class with an __enter__() and an __exit__() method can work as a context manager.
Double underscore methods are special methods that are not called but instead triggered. They are internally set to run at specific times or after certain events. For example, __enter()__ runs when a with statement is entered and __exit__ runs right before the with block is left.
We have seen an example of function-based context managers, which work great for quick and simple cases. For more complex use cases, we'll define context management as an additional ability to an existing class.
Let's create a ListProtect context manager class, which will operate on a copy of a list before returning the changes. This way, the original list would be restored unaltered if an error occurred during the operation.
Let's see how that looks.
class ListProtect:
def __init__(self, original: list):
self.original = original
def __enter__(self):
self.clone: list = self.original.copy()
return self.clone
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is None:
self.original[:] = self.clone
else:
print('@ListProtect: Error occurred while processing the list. The changes are discarded.')
return TrueAnd we can use it like so:
the_list = [1, 2, 3]
# Error-free example
with ListProtect(the_list) as the_copy:
the_copy.append(100)
print(f'the list: {the_list}')The list was successfully modified. But what if an error occurs?
the_list = [1, 2, 3]
# Example with errors
with ListProtect(the_list) as the_copy:
the_copy.append(100)
the_copy.append(1/0)
print(f'the list: {the_list}')In this instance, ListProtect caught an error and the list remained unaltered. So how does this work?
Explanation
In the ListProtect class, we defined two required methods:
__enter__() - this method defines what happens before the logic under the with statement runs. If the enter method returns anything, it is bound to the context variable. In our class, we create a clone of the list and returned it as the context variable—the_copy in the above examples.
__exit__() - this method defines what happens when the with logic is complete or has raised an error. Besides self, this method takes three parameters: exc_type, exc_val, exc_tb, which can also be shortened to *exc. If no exception occurs in the managee, all these values are None. If we return a truthy value at the end of __exit()__—as we did in our class—the error is suppressed. Otherwise, if we return a falsy value, such as None, False, or blank, the error is propagated.
In short, ListProtect's __exit()__ method first checks whether the exception type was None. If so, there were no errors, and it applies the changes to the original list; otherwise, it announces an error occurred.
Summary
A context manager, used in a with statement, defines a temporary context for the given set of operations. It does so by injecting code at the beginning and at the end of the code block, setting up the context at the beginning, and tearing it down at the end.
Meet the Authors
Founder of LearnDataSci