Python: context manager

Or how to prepare things before working with them, and then ensure they are correctly released.

A common case of using a context manager.

What is a context manager?

It is something you often use without knowing how it does work!

It is associated to the with statement, you can see an example below where a file is read.

with open('file.json', encoding="utf-8") as fp:
    config = json.loads(fp.read())

What does it do?

It ensures that things are available only inside the context.

When entering a context we can automatically perform some actions like:

• opening a file, or a database connection,
• obtaining a lock,
• temporarily change a setting/configuration,
• …

When exiting the context, we would then:

• close the file or database connection,
• release the lock,
• go back to the previous setting/configuration value
• …

In practice

In the case of the open function, it does not do much when entering the context, but it does something when exiting the context: it closes the file.

Without using a context manager you would have had to write something like:

fp = open("file.json", encoding="utf-8")
try:
    config = json.loads(fp.read())
finally:
    fp.close()  # always close the fp at the end of the try block

(As the file should always be closed, the close method should be put on a finally block)

Please notice that the context manager is associated with a value and not with a function/method, the following example would work as well as the first one (but it would be less elegant).

fp = open('file.json', encoding="utf-8")
with fp:
    config = json.loads(fp.read())

How to create a custom one?

There are two way to create a custom context manager, the first is decorating a function, and the second one is by implementing the __enter__ and __exit__ method in a class (or __aenter__ and __aexit__ in the case of an async context manager)

Creating a context manager using a decorator

The source code is taken from the official python documentation.

# First creating the context manager
from contextlib import contextmanager

@contextmanager
def managed_resource(*args, **kwds):
    # Code to acquire resource, e.g.:
    resource = acquire_resource(*args, **kwds)
    try:
        yield resource
    finally:
        # Code to release resource, e.g.:
        release_resource(resource)

# Then using it
with managed_resource(timeout=3600) as resource:
    # Resource is released at the end of this block,
    # even if code in the block raises an exception

The yield resource in the function is equivalent to a return, except that it does not stop the execution of the function.

The finally block is called only once the resource isn’t used anymore: when leaving the with block.

Allowing a class instance to be used as a context manager

The following code is the equivalent from the previous one

# First creating the context manager
class ManagedResource:
    def __init__(self, *args, **kwargs):
        self._args = args
        self._kwargs = kwargs

    def __enter__(self):
        self.resource = acquire_resource(*self._args, **.self._kwargs)

    def __exit__(self, exc_type, exc_val, exc_tb):
        release_resource(self.resource)

# Then using it
with ManagedResource(timeout=3600) as resource:
    # Resource is released at the end of this block,
    # even if code in the block raises an exception

On the __exit__ method, exc_type, exc_value and exc_tb parameters allow to have access to an exception that could be raised inside the context. Please notice that the __exit__ method will always be executed.

When handling an exception inside the __exit__ method, it should return True if the exception was fully handled, any other value would mean that the exception would need to be handled in the code surrounding the with statement.

Example of using a custom context manager

The following class extract is based on how should work Apache Airflow DAG implementation, it is taken from a site project.

class TaskDAG:

    CURRENT: List["TaskDAG"] = []

    @classmethod
    def get_dag(cls) -> Optional["TaskDAG"]:
        return cls.CURRENT[-1] if len(cls.CURRENT) else None

    def __enter__(self) -> Self:
        self.__class__.CURRENT.append(self)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.__class__.CURRENT.pop()
        from core.context.global_context import GlobalContext

        GlobalContext.get_instance().register_dag(self)

When inside a context, a task can register itself to the latest stored TaskDAG.

class Task:
    def __init__(self, dag: TaskDAG = None):
        current_dag = dag if dag else TaskDAG.CURRENT[-1]  # last DAG in the context

        current_dag.register_task(self)

# using a context manager
with TaskDAG():
    # no need to pass dag value, it will be retrieved
    task = Task()

# replace the following syntax
dag = TaskDAG()
task = Task(dag=dag)

We will agree that the context manager syntax is more elegant than the classical one.

That’s all folks!