Python Context Managers Clearly Explained

Introduction

Python context managers provide a convenient way to manage resources, such as files, network connections, and locks. They enable automatic setup and teardown operations, saving development effort, simplifying code, minimizing errors. Context managers further ensure that resources are properly managed even in the presence of exceptions or errors.

Understanding Python Context Managers

A context manager is an object that defines a runtime context. This means that the object runs some predefined code when entering the context, provides an object to operate on within the context, and finally runs some other predefined code when exiting the context.

For example, the context can be operations on a file. The context manager opens the file when entering the context, provides a file object, and closes it upon exit.

In Python, context managers are objects that implement a pair of methods: __enter__() and __exit__(). The __enter__() method is called when entering the context. It runs any necessary setup code and returns the context object. (Returning a context object is in fact optional, and if not done, the actual returned object is None, as is the case for any other Python function that does not return a value.) The __exit__() method is called when exiting the context, and it can perform cleanup operations, such as releasing resources or handling exceptions.

Using Context Managers with the Python with Statement

The Python with statement provides the syntax for using context managers. It automatically calls the __enter__() method when entering its code block, i.e. the context, and ensures that the __exit__() method is called when exiting, even if an exception occurs within the block.

The following example demonstrates using the with statement to open and read a text file. (For the sake of the example’s output, the file is assumed to contain two comma-separated integers, “3, 5”). It then prints out the resulting decimal number upon division.

Example:

with open("fraction.txt", "r") as file:
    content = file.read().strip()
    numerator, denominator = map(int, content.split(','))
    print(numerator/denominator)

Output:

0.6

In the example above, the open() function opens the file and returns the context manager (line 1). The with statement calls the context manager’s __enter__() method and assigns the object it returns to the file variable. The body of the with statement uses the file variable to read the file’s text (line 2), convert the content to two integers (line 3), and print the decimal (line 4).

Finally, the with statement calls the context manager’s __exit__() method, which closes the file. The with statement calls __exit__() covertly for us, and in this case, we can check that the file was indeed closed by calling file.read() again. When we do, we get the following error: “ValueError: I/O operation on closed file.”

The with statement guarantees that the __enter__() function runs regardless of any exception within its code block. For instance, in the example above, if the file contains text that can’t be converted to integers or a denominator whose value is 0, an exception occurs, but the file still closes.

This is similar to the way the finally clause of a try–except–finally exception handling statement works. But when using with, you don’t have to explicitly write a finally clause, or even know what teardown code you should run. The context manager handles all that for you.

Here is the basic syntax of the with statement:

with context_manager as obj:
   ... # The context code block

These are the key elements:

context_manager: The context manager that implements the __enter__() and __exit__() methods. The statement will call __enter__() at the start of the block and __exit__() at the end.
obj: The with statement automatically assigns the object that the __enter__() method returns, i.e. the context object, to this variable. The as obj part of the syntax can be omitted when the code block does not need to use the context object, as you’ll see in the timer example below.

Creating Custom Context Managers Using a Class

While Python provides built-in context managers for common tasks such as file handling, you can also create custom context managers for your specific needs. You can define a custom context manager by implementing a class with __enter__() and __exit__() methods.

Here is an example of a context manager that times the execution of a block of code:

Example:

import time

class Timer:
    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        execution_time = self.end_time - self.start_time
        print(f"Execution time: {execution_time} seconds")

In the example above, the Timer class is a context manager (line 3). It defines the __enter__() (line 2) method, which records the start time (line 5). It also defines the __exit__() method (line 8), which calculates the execution time by subtracting the start time from the end time and then prints it (lines 9-11).

As you can see, aside from the self parameter, which any instance method must have, __exit__() also has the following exception-handling-related parameters:

exc_type: The type of exception that occurred, if any. If no exception occurs, its value is None.
exc_value: The exception instance that was raised, if any. If no exception occurs, its value is None.
traceback: A traceback object representing the call stack at the point where the exception occurred, if any. If no exception occurs, its value is None.

These parameters allow the __exit__() method to handle any exceptions that occur within the context and perform appropriate cleanup operations. If the method returns True, it indicates that any exception that occurred has been handled, and the exception won’t be propagated further. If it returns False or nothing (implicitly None), the exception will be propagated as usual after the __exit__() method completes its execution.

You can use the Timer context manager with the with statement to time any block of code:

Example:

with Timer():
    time.sleep(2)  # Code block to be timed

Output:

Execution time: 2.0051629543304443 seconds

Creating Custom Context Managers With a Generator Function Using _contextlib_

The contextlib module also provides a convenient way for creating context managers out of generator functions using the contextmanager decorator. This simplifies the process of creating generators because instead of defining a class with the two __enter__() and __exit__() methods, all you have to do is define just one function.

A generator function executes code until it encounters a yield statement, at which point it stops. When it’s called again, it picks up right after the yield and continues executing its code from there.

Using a generator function and the contextmanager decorator, the following example creates a custom timer context manager similar to the one implemented in the previous example.

Example:

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start_time = time.time()
    yield
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"Execution time: {execution_time} seconds")

In the example above, the timer() generator function is decorated with @contextmanager (line 4), transforming it into a function that returns a context manager object. Inside the function, the start time is recorded (line 6) before yielding control to the code block within the with statement (line 7). After the block completes, the end time is recorded (line 8), and the execution time is calculated (line 9) and printed (line 10). The yield statement separates the code that executes before entering the context from when exiting it.

The usage of this type of context manager is the same as that of a class-based context manager.

Example:

with timer():
    time.sleep(2)  # Code block to be timed

Output:

Execution time: 2.002261161804199 seconds

In example, the timer() function call returns a generator, while in the previous example, we had Timer(), which created a new object of the Timer class.

Here is the basic code for using a generator function-based context manager:

from contextlib import contextmanager

@contextmanager
def generator_fn():
    ... # Code that executes when entering the context
    yield context_object
     ... # Code that executes when exiting the context

These are the key elements:

from contextlib import contextmanager: Imports the contextmanager decorator so that it’s available for use.
generator_fn(): The generator function to be used as a context manager. You can also define it with any parameters you wish to change the behavior of the context manager.
yield: This yield is the end of the context manager’s entry code, i.e. the code that’s in __enter__(). Following the yield is the code that __exit__() runs.
context_object: This is the object that the yield statement yields. It becomes the object that __enter__() returns. It is optional, and if not given, __enter__() returns None.

Understanding How the _contextmanager_ Decorator Works

The contextmanager decorator uses the generator function to create the __enter__() and __exit__() methods of a generator class. It then returns a function that returns that class.

Let’s see how this works by implementing our own version of contextmanager:

def my_contextmanager(fn):
    class GeneratorContextManager():
        def __init__(self, *args, **kwargs):
            self.generator = fn(*args, **kwargs)

        def __enter__(self):
            return next(self.generator)

        def __exit__(self, exc_type, exc_value, traceback):
            try:
                next(self.generator)
            except StopIteration:
                pass

    def generator_context_manager(*args, **kwargs):
        return GeneratorContextManager(*args, **kwargs)

    return generator_context_manager

The above code defines the my_contextmanager() decorator function (line 1), which is the “homemade” version of contextmanager. This function defines an inner class, GeneratorContextManager (line 2), and an inner function, generator_context_manager() (line 15).

All the interesting parts happen in the body of the GeneratorContextManager class. The generator_context_manager() function merely returns the class (line 16). It’s there so that the decorator can comply with the essence of being a function decorator, which is a function that returns another function.

Within the GeneratorContextManager class, the initializer, __init__() (line 3), stores the generator function in the instance variable generator (line 4). The generator function comes from the fn parameter of the decorator (line 1). (It’s the function the decorator is decorating.) The generator variable retains the generator function within the class for use by the __enter__() and __exit__() methods that implement the context manager protocol.

The __enter__() method calls next() on the generator and returns the resulting value (line 7). This executes all the code in the generator up to its yield statement and returns the yielded value.

The __exit__() method calls next() on the generator again, which will execute the code that comes after the generator’s yield statement (line 11). Since the generator has only one yield and completes its last iteration at this point, this next() call generates a StopIteration exception. In order to prevent this exception from interrupting execution of code, __exit__() handles it with a try–catch statement (lines 10-13).

If you understand this code, then congratulations! You now have a good grasp of context managers, as well as generators and decorators.

The following example tests the my_contextmanager decorator by using it. It also demonstrates a context manager that has an argument.

Example:

@my_contextmanager
def conversation_wrapper(name):
    print(f"Hello, {name}.")
    yield name
    print(f"Goodbye, {name}.")

with conversation_wrapper("Snoopy") as name:
    print("Glad to meet you!")
    print(f"{name}, are you enjoying learning Python?")
    print("It's fantastic!")

Output:

Hello, Snoopy.
Glad to meet you!
Snoopy, are you enjoying learning Python?
It's fantastic!
Goodbye, Snoopy.

The conversation_wrapper() context manager in the example (lines 2-5) above is capable of wrapping a printed conversation with a greeting and a goodbye message, both containing a name given as an argument. The body of the conversation can also include the name, and therefore, it’s yielded and made available as a resource in the context (line 4). The body of the conversation is printed within the with statement (lines 7-10).

Summary & Reference for Python Context Managers

Python context managers provide a convenient way to manage resources, such as files, network connections, and locks. They simplify resource management by automating setup and teardown operations and ensuring proper resource handling, even when exceptions occur.

Context managers are objects that implement a pair of methods: __enter__() and __exit__(). The __enter__() method is called when entering the context, and it optionally returns a context object. The __exit__() method is called when exiting the context.

The Python with statement provides the syntax for using context managers. It automatically calls the __enter__() method when entering its code block, i.e. context, and ensures that the __exit__() method is called when exiting. It also assigns the context object to a local variable that follows the as keyword. The as keyword is optional and can be omitted when the context object is not needed within the code block.

with context_manager as obj:
   ... # The context code block

Custom context managers can be created using either classes or generator functions.

To create a context manager with a class, implement a class with __enter__() and __exit__() methods.

class Timer:
    def __enter__(self)
        ... # Context entry code (setup)
        return context_obj  # Optional

    def __exit__(self, exc_type, exc_value, traceback):
        ... # Context exit code (teardown)
        return True  # Optionally return True to suppress any exceptions in the context that have been handled by this method

To create a context manager with a generator function, decorate it with the contextmanager decorator of the contextlib module. Place a single yield statement within the function to separate the code executed when the context is entered from when it’s exited. The value yielded by yield becomes the context object.

from contextlib import contextmanager

@contextmanager
def generator_fn():
    ... # Code that executes when entering the context
    yield context_object
     ... # Code that executes when exiting the context