| jupytext |
|
||||||||
|---|---|---|---|---|---|---|---|---|---|
| kernelspec |
|
As you are (hopefully) already accustomed to (at least basic) Python, you've almost certainly encountered those fancy and extremely useful structures with the {code}with keyword, called context managers. Where does this name come from? Well, context managers define a runtime environment (context) for a given block of code.
:class: tip
I am sure you know it, but a block of code is a sequence of Python statements within the same indentation, e.g., this block of code for the `for` loop contains just a single statement: {code}`print(i)`:
```python
for i in range(10):
print(i)
```
Why do we need such an environment? Well, when you need to configure anything before executing some piece of code (a code block) and/or tear down and clean up afterwards. That's the core idea behind context managers: set up before, clean up after, guaranteed.
To run code within the context managed by a context manager, we use the with statement. If the context manager returns any object on setup, you can bind it using the {code}as keyword:
:tags: [remove-stderr]
with open("some_dummy_file.txt", "w") as f:
f.write("1st line\n")
f.write("2nd line")
Context managers can be created in two ways: either as a class-based definition or a generator-based definition. Both are presented below.
:class: important
To define a context manager using classes, you need to know what a **class** is and how we can define it. If you don't have such knowledge, learn about it and return later.
A context manager can be created as a class which implements two special (dunder1) methods:
__enter____exit__
It is quite clear that the first (__enter__) is responsible for setting up the execution context (environment), whereas the latter is responsible for cleaning up.
The setup method is argumentless, namely it takes just an implicit parameter (self), but you can return anything. Remember that the object returned by the __enter__ method can be bound to the target variable using the as keyword.
:lineno-start: 1
:emphasize-lines: 5
class MyContextManager:
def __enter__(self):
# some logic goes here
return self # We return the object itself to be able to bind it using `as`
:class: tip
If the `__enter__` method does not raise an error, the `__exit__` method is guaranteed to be invoked {cite}`python:compound_with`, regardless of whether the block inside the context manager raises an error or not.
Unlike __enter__, the __exit__ method imposes three parameters:
- exception type (type hint:
type[Exception] | None), - exception value (type hint:
Exception | None), - traceback (type hint:
types.TracebackType | None)
If the block of code inside the context manager does not raise an error, all those arguments will be None. Otherwise, they will be set accordingly. The question is, what to do next with such an error? You have two options:
- suppress it (silence it, do not propagate),
- reraise it
To suppress the exception (or manage it on your own inside the __exit__ method), you need to return the value True explicitly or any expression evaluated to True {cite}python:stdtypes_truth. Otherwise, the exception will be propagated. Following our above example, let us add a dummy __exit__ implementation:
:lineno-start: 7
:emphasize-lines: 3
def __exit__(self, exc_type: type[Exception] | None, exc_val: Exception | None, traceback):
# some logic goes here
return True # If we want to suppress the error, we return True
As an example, we will try to reimplement a simplified version of the chdir context manager from contextlib standard library to change the current working directory.
:lineno-start: 1
:emphasize-lines: 7, 10, 11, 16
import os
class ChDir:
saved_path: str
new_path: str
def __init__(self, new_path):
self.new_path = new_path # We save the new directory we move into inside the context manager
def __enter__(self):
self.saved_path = os.getcwd() # We store this to restore it later
os.chdir(self.new_path) # We actually change directory here
print(f"📂 Changed directory to: {self.new_path}")
return self # Not required, but we can use it with the `as` keyword
def __exit__(self, exc_type, exc_value, traceback):
os.chdir(self.saved_path) # We restore the original working directory
print(f"↩️ Returned to: {self.saved_path}")
:class: attention
To create asynchronous context manager (see section below [Asynchronous Context Managers](#asynchronous-context-managers)) {cite}`python:datamodel_asynccontext`, you need to use [`__aenter__`](https://docs.python.org/3/reference/datamodel.html#object.__aenter__) and [`__aexit__`](https://docs.python.org/3/reference/datamodel.html#object.__aexit__) asynchronous dunder methods. Their signature is the same as their synchronous counterparts.
:class: important
Before you dive into this section, ensure you are familiar with Python generators. Read 📰 about them in the chapter [Generators](./generators.md).
:class: important
You should be aware of decorators before reading this section. You will find them in the chapter [Decorators](./decorators.md).
Writing context managers as classes gives us the most flexibility; however, there is another, quite convenient way to create a simple context manager—by using a generator function {cite}python:generator_glossary_misc. The pivotal element is the yield statement and the @contextmanager decorator from the contextlib {cite}python:contextmanager_doc module, which is used to decorate the generator function. It has the following structure:
:lineno-start: 1
:emphasize-lines: 3, 7
from contextlib import contextmanager
@contextmanager # We need to use this decorator to make a func a context manager
def my_context_manager():
# here is the logic to run on setup
try:
yield # We can yield something to be able to bind it using `as`
except Exception as e:
# logic to handle (optionally) the exception
finally:
# logic to clean up
So, we have:
:header-rows: 1
:name: code-explanation-context-manager
* - **Line**
- **Code**
- **Explanation**
* - 1
- ```python
from contextlib import contextmanager
```
- Imports the `contextmanager` decorator from the `contextlib` module of the standard Python library (no need to install anything)
* - 3-4
- ```python
@contextmanager
def my_context_manager():
```
- We decorate the function to make it a context manager
* - 5
- ```python
# here is the logic to run on setup
```
- Here (before the `yield` statement) you put all the logic to set up the runtime context (as you would in the `__enter__` method of a class-based definition)
* - 6-7
- ```python
try:
yield
```
- This is where the magic happens. It creates a generator function, so it saves the state and sends a value (here we send nothing, but anything you add after `yield` can be assigned using the `as` keyword, like the return value of the `__enter__` method). It is inside the `try-except` block to be able to catch and process (or suppress) any exception that occurred in the context manager
* - 8-9
- ```python
except Exception as e:
# logic to handle (optionally) exception
```
- Here we catch any exception which occurred inside the context manager. You can catch a more specific error or keep it as general as `Exception`. Here you have flexibility to handle or suppress a particular type of exception
* - 10-11
- ```python
finally:
# logic to clean up
```
- As in most programming languages, the `finally` block is executed in either case: whether an error was raised or not. The logic inside the `finally` block corresponds to the cleanup logic you would put in the `__exit__` method of a class-based definition
:class: attention
If you intend to create an asynchronous context manager (see section below [Asynchronous Context Managers](#asynchronous-context-managers)), remember to define the function (decorated with [`@asynccontextmanager`](https://docs.python.org/3/library/contextlib.html#contextlib.asynccontextmanager)) as `async`.
:class: tip
The [`contextlib`](https://docs.python.org/3/library/contextlib.html) library provides a variety of other useful tools. For example, you can create a context manager **as a decorator**. To do that, you just follow the [class-based definition](#class-based-definition) and add a parent class [`ContextDecorator`](https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator), so you can use it as:
```{code-block} python
:lineno-start: 1
:emphasize-lines: 4
import os
from contextlib import ContextDecorator
class ChDir(ContextDecorator): # Note, we've added a parent class
...
@ChDir("/tmp")
def some_func():
# some function's logic
print(f"I am in {os.getcwd()}")
some_func()
```
Regardless of the way you defined your context manager, you can use it as below:
:tags: ["remove-input", "remove-output"]
import os
class ChDir:
saved_path: str
def __init__(self, new_path):
self.new_path = new_path # We save the new directory we move into inside the context manager
def __enter__(self):
self.saved_path = os.getcwd() # We store this to restore it later
os.chdir(self.new_path) # We actually change directory here
print(f"📂 Changed directory to: {self.new_path}")
return self # Not required, but we can use it with the `as` keyword
def __exit__(self, exc_type, exc_value, traceback):
os.chdir(self.saved_path) # We restore the original working directory
print(f"↩️ Returned to: {self.saved_path}")
os.chdir("/home")
print("Before context:", os.getcwd())
with ChDir("/tmp"): # or any folder you have
print("Inside context:", os.getcwd())
# do stuff here safely
print("After context:", os.getcwd())
:class: important
To understand this section well, read 📰 the [Coroutines](./coroutines.md) chapter first.
When possible, such as when performing input/output operations like handling files, opening database connections, or managing HTTP connections, we can create an asynchronous version of a context manager. The dunder methods for creating asynchronous context managers can be coroutines (see the Coroutines chapter). These special methods' names change slightly: they must be defined using async def, and their names are __aenter__ and __aexit__ for setting up and cleaning up the asynchronous context manager, respectively. As an example, let's create a simple asynchronous context manager to lock a file in order to prevent it from being overwritten by concurrent tasks (to avoid data inconsistency).
:lineno-start: 1
:emphasize-lines: 11, 17
import asyncio
class AsyncFileLock:
filename: str
lock: asyncio.Lock
def __init__(self, filename):
self.filename = filename
self.lock = asyncio.Lock()
async def __aenter__(self): # Note: we are using `async` and the name changed
print(f"Acquiring lock for {self.filename}...")
await self.lock.acquire()
print(f"Lock acquired for {self.filename}")
return self
async def __aexit__(self, exc_type, exc_val, exc_tb): # Note: the name is `__aexit__`
self.lock.release()
print(f"Lock released for {self.filename}")
return False
We can then use it as follows:
:tags: ["remove-input", "remove-output"]
import asyncio
class Lock:
filename: str
lock: asyncio.Lock
def __init__(self, filename):
self.filename = filename
self.lock = asyncio.Lock()
async def __aenter__(self):
print(f"Acquiring lock for {self.filename}...")
await self.lock.acquire()
print(f"Lock acquired for {self.filename}")
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
self.lock.release()
print(f"Lock released for {self.filename}")
return False
async def some_long_file_handling_func():
print("Processing file...")
await asyncio.sleep(10)
print("File processing done")
async with Lock("/tmp/sample.txt"):
await some_long_file_handling_func()
Footnotes
-
Don't remember what dunder methods are? See Special Attributes. ↩