Login

  • Jan 28, 2025

Custom Context Manager in Python

Context manager is useful when dealing with resources. Python has many built-in context managers. However, it lets you define your own context manager. This article presents how to do this.

Introduction

The with statement in Python, assists you in implementing some popular resource management patterns by abstracting their functionality, allowing them to be factored out and reused.

The with statement is commonly used for file handling, database connections, locks management in multi-threaded environment and other context-based operations.

The with statement improves readability.

In other words, you can say that in Python, with statement is used in exception handling but with the cleaner and readable code.

The Python with statement creates a runtime context that allows you to run a group of statements under the control of a context manager.

The syntax

The with statement uses the following general syntax:

with expression as target_var:
    do_process(target_var)

The Context Manager and Context Manager Protocol

Look at the syntax, The context manager object results from evaluating the expression after the with statement. In other words, the expression must return an object that implements the context manager protocol. This protocol consists of two special methods:

1 — __enter__() is called by the with statement to enter the runtime context.

2 — __exit__() is called when the execution leaves the “with” code block.

In the syntax, the as specifier is optional. If you provide a target_var using as keyword, then the return value of calling .__enter__() on the context manager object is bound to that variable.

Steps of context manager protocol

Here’s how the with statement proceeds when Python runs into it:

1 — Call expression to get a context manager.

2 — Store the context manager’s .__enter__() and .__exit__() methods for later use.

3 — Invokes “.__enter__()” in the context manager. Bind the returned value which is received from “.__enter__()” to the target_varif any value is returned.

Some context managers return None from .__enter__() because they have no useful object to give back to the caller. In such cases, using the target_var is not required, hence may be omitted.

4 — Execute the with code block i.e. all the statement within the with statement.

5 — Invoke .__exit__() method on the context manager when the exiting the with statement i.e. when execution of withcode block is over.

The .__enter__() usually contains the setup code. Like a conditional statement or a for loop, the with statement is also a compound statement. Hence, the with statement initiates the execution of a code block i.e. to run many statements. Most often, you utilize the with code block to deal with target_var if any value is returned by the .__enter__()method.

Once the ‌code block is completed, the .__exit__() method is called. This method usually contains the deconstruction logic or cleanup code, such as calling the  .close() method on an open file object during file processing. That is why the with statement is so popular since it is very effective at appropriately acquiring and releasing resources.

Creating Custom Context Managers

You already know that with statement works with a context manager which is obtained by executing the expression which follows the keyword with. The open() function for working with file objects, threading.Lock() for working with threads, decimal.localcontext() for working with high precision computation etc. returns a context manager.

You may achieve the functionality of a context manager by defining a custom context manager.

There are two ways to implement custom context manager (context management protocol):

1 — Using classes

2 — Using functions

To implement the functionality of a context manager, you need to implement the context management protocol.

Implementation of a context management protocol means you need to implement both special methods —  .__enter__()and .__exit__() in your class-based context managers.

In general, context managers and the ‌statement aren’t just for resource management. They enable you to share and reuse common setup and teardown (cleanup) code. In other words, context managers allow you to do any pair of actions that must be done before and after another operation or procedure, such as:

  • Open and close — while file handling

  • Lock and release — in case of multithreading.

  • Change and reset — in case of rollback.

  • Create and delete.

  • Enter and exit — in case of mutual exclusion.

  • Start and stop.

  • Setup and teardown — in case of database connection

Any of the above pairs of operations can be safely managed using code in a context manager. Then you can use that context manager in statements throughout your code. This eliminates errors and reduces boilerplate code.

You can create your own context managers by defining classes or functions. Let’s explore the first approach in the next section.

Class-Based Context Manager

You can create your own context managers by defining classes with __enter__() and __exit__() methods. This ‌can be demonstrated with the help of the below examples.

Example 1:

Code:

class CustomContextManager:
    def __init__(self, file_name):
        self.file_name = file_name
    def __enter__(self):
        self.file = open(self.file_name, 'w')
        return self.file
    def __exit__(self, *args):
        self.file.close()
#usage of the above class
with CustomContextManager ('my_file.txt') as file:
    file.write ('Hello! I am a custom class-based context manager!')

Output:

The above code will create a file — my_file.txt in the same directory as the above code file and the message in the file. The output in my system is shown in the below screenshot which also depicts my code file “demo.py”.

Custom Class Based Context Manager

In this example, the __enter__() method initializes the resource (file) and returns a descriptor. The __exit__() method ensures proper cleanup.

Example 2: Changing standard output

The below defined context manager receives a file object via its constructor. In .__enter__(), you first save the standard output — sys.stdout, to the instance attribute —  original_streamsince you need to restore to this original stream when the with block (context) is left. Then you reconfigure the standard output to point to the file object saved in the other instance attribute target_stream which is the received file object. In.__exit__(), you just restore the standard output to its original value.

Further, the code line

with CustomContextManagerRedirectStdout(f):

has no asspecifier since the above context manager — CustomContextManagerRedirectStdout is an example of a context manager which does not return a value from the .__enter__() method unlike the context manager in example-1, the  .__enter__() method returned an open file object.

Code

import sys
class CustomContextManagerRedirectStdout:
    def __init__(self, target_stream):
        self.target_stream = target_stream
    def __enter__(self):
        # save the original standard output stream
        self.original_stream = sys.stdout
        # change the original standard output stream to the target_stream i.e. to the file-name received
        sys.stdout = self.target_stream
    def __exit__(self, exc_type, exc_value, traceback):
         # change the original standard output stream to the target_stream i.e. to the file-name received
        sys.stdout = self.original_stream
# Usage:
with open("output.txt", "w") as f:
    with CustomContextManagerRedirectStdout(f):
        print ("Standard output is set to the file 'output.txt' with the help of a custom context manager.")
        print ("Hence, all the output of print statement will go to the file 'output.txt'.")
        print ("After the exit from the with block i.e. outside the with block, the standard output is reset back to console")
        
# Outside the 'with' block, stdout is restored to its original state
print ("The standard output is reset back to standard console output - the original stream")

Output:

The code will produce two outputs — (1) the output of the three print statements within the with statement written to file “output.txt” since the standard is redirected to this file. (2) the output of the last print statement outside with statement goes to console output. These outputs are shown in the screenshots below.

First Output of custom context manager

Second Output of custom context manager

A question is still unanswered from the above code snippet. The signature of the __ext__() method is — __exit__(self, exc_type, exc_value, traceback)

What about the last three arguments in the method. When the flow of execution exits the context, the .__exit__()method is invoked. If there was none exception occurred inside the with block, the three last arguments of .__exit__()are set to None. Otherwise, they contain the type, value, and traceback associated with the exception in question. More on this in a different article

Recapitulate

This article presents how to implement a custom class-based context manager. In ‌another article, I will come up with a custom function-based context manager.

Visit our website at https://www.devtechie.com for a wealth of additional learning resources.