Sachin

Posted on Feb 27 • Originally published at geekpython.in

Python: Guide to Context Manager & How do they work?

#python #programming

In this article, we'll look at context managers and how they can be used with Python's "with" statements and how to create our own custom context manager.

What Is Context Manager?

Resource management is critical in any programming language, and the use of system resources in programs is common.

Assume we are working on a project where we need to establish a database connection or perform file operations; these operations consume resources that are limited in supply, so they must be released after use; otherwise, issues such as running out of memory or file descriptors, or exceeding the maximum number of connections or network bandwidth can arise.

Context managers come to the rescue in these situations; they are used to prepare resources for use by the program and then free resources when the resources are no longer required, even if exceptions have occurred.

Why Use Context Manager?

As previously discussed, context managers provide a mechanism for the setup and teardown of the resources associated with the program. It improves the readability, conciseness, and maintainability of the code.

Consider the following example, in which we perform a file writing operation without using the with statement.

# Opening file
file = open('sample.txt', 'w')
try:
    # Writing data into file
    data = file.write("Hello")
except Exception as e:
    print(f"Error Occurred: {e}")
finally:
    # Closing the file
    file.close()

To begin, we had to write more lines of code in this approach, and we had to manually close the file in the finally block.

Even if an exception occurs, finally block will ensure that the file is closed. However, using the open() function with the with statement reduces the excess code and eliminates the need to manually close the file.

with open("sample.txt", "w") as file:
    data = file.write("Hello")

In the preceding code, when the with statement is executed, the open() function's __enter__ method is called, which returns a file object. The file object is then assigned to the variable file by the as clause, and the content of the sample.txt file is written using the variable file. Finally, when the program exits execution, the __exit__ method is invoked to close the file.

We'll learn more about __enter__ and __exit__ methods in the upcoming sections.

We can check if the file is actually closed or not.

print(file.closed)

----------
True

We received the result True, indicating that the file is automatically closed once the execution exits the with block.

Using with Statement

If you used the with statement, it is likely that you also used the context manager. The with statement is probably most commonly used when opening a file.

# Opening a file
with open('sample.txt', 'r') as file:
    content = file.read()

Here's a simple program that opens a text file and reads the content. When the open() function is evaluated after the with statement, context manager is obtained.

The context manager implements two methods called __enter__ and __exit__. The __enter__ method is called at the start to prepare the resource to be used, and the __exit__ method is called at the end to release resources.

Python runs the above code in the following order:

The with statement is executed, and the open() function is called.
The open() function's __enter__ method opens the file and returns the file object. The as clause then assigns the file object to the file variable.
The inner block of the code content = file.read() gets executed.
In the end, the __exit__ method is called to perform the cleanup and closing of the file.

Let's define and implement both these methods in a Python class and try to understand the execution flow of the program.

Creating Context Manager

The context manager will be created by implementing the __enter__ and __exit__ methods within the class. Any class that has both of these methods can act as a context manager.

Defining a Python class

# Creating a class-based context manager
class Conmanager:
    def __enter__(self):
        print("Context Manager's enter method is called.")

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Exit method is called...")
        print(f'Exception Type: {exc_type}')
        print(f'Exception Value: {exc_val}')
        print(f'Exception Traceback: {exc_tb}')

# Using the "with" stmt
with Conmanager() as cn:
    print("Inner block of code within the 'with' statement.")

First, we created a class named Conmanager and defined the __enter__ and __exit__ methods inside the class. Then we created the Conmanager object and assigned it to the variable cn using the as clause. We will get the following output after running the above program.

Context Manager's enter method is called.
Inner block of code within the 'with' statement.
Exit method is called...
Exception Type: None
Exception Value: None
Exception Traceback: None

When the with block is executed, Python orders the execution flow as follows:

As we can see from the output, the __enter__ method is called first.
The code contained within the with statement is executed.
To exit the with statement block, the __exit__ method is called at the end.

We can see in the output that we got None values for the exc_type, exc_val, and exc_tb parameters passed inside the __exit__ method of the class Conmanager.

When an exception occurs while executing the with statement, these parameters take effect.

exc_type - displays the type of exception.
exc_val - displays the message of the exception.
exc_tb - displays the traceback object of the exception.

Consider the following example, which shows how these parameters were used when an exception occurred.

# Creating a class-based context manager
class Conmanager:
    def __enter__(self):
        print("Enter method is called.")

        return "Do some stuff"

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Exit method is called...")
        print(f'Exception Type: {exc_type}')
        print(f'Exception Value: {exc_val}')
        print(f'Exception Traceback: {exc_tb}')

# Using the "with" stmt
with Conmanager() as cn:
    print(cn)
    # Raising exception on purpose
    cn.read()

When we run the above code, we get the following result.

Enter method is called.
Do some stuff
Exit method is called...
Exception Type: <class 'AttributeError'>
Exception Value: 'str' object has no attribute 'read'
Exception Traceback: <traceback object at 0x00000276057D4800>
Traceback (most recent call last):
  ....
    cn.read()
AttributeError: 'str' object has no attribute 'read'

Instead of getting None values, we got the AttributeError, as shown in the output above and those three parameters displayed certain values.

exc_type displayed the <class 'AttributeError'> value.
exc_val displayed the 'str' object has no attribute 'read' message.
exc_tb displayed the <traceback object at 0x00000276057D4800> value.

Example

In the following example, we've created a context manager class that will reverse a sequence.

class Reverse:
    def __init__(self, data):
        self.data = data

    def __enter__(self):
        self.operate = self.data[:: -1]
        return self.operate

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass


with Reverse("Geek") as rev:
    print(f"Reversed string: {rev}")

We've created a class called Reverse and defined the __init__ method, which takes data, the __enter__ method, which operates on the data and returns the reversed version of it, and the __exit__ method, which does nothing.

Then we used the with statement to call the context manager's object, passing the sequence "Geek" and assigning it to the rev using the as clause before printing it. We will get the following output after running the above code.

Reversed string: keeG

The upper code contains a flaw because we did not include any exception-handling code within the __exit__ method. What if we run into an exception?

with Reverse("Geek") as rev:
    # Modified the code from here
    print(rev.copy())

We changed the code within the with statement and attempted to print the rev.copy(). This will result in an error.

Traceback (most recent call last):
  ....
    print(f"Reversed string: {rev.copy()}")
AttributeError: 'str' object has no attribute 'copy'

Exception Handling

Let's include the exception handling code in the __exit__ method.

class Reverse:
    def __init__(self, data):
        self.data = data

    def __enter__(self):
        self.operate = self.data[:: -1]
        return self.operate

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            return self.operate
        else:
            print(f"Exception occurred: {exc_type}")
            print(f"Exception message: {exc_val}")
            return True

with Reverse("Geek") as rev:
    print(rev.copy())

print("Execution of the program continues...")

First, we defined the condition to return the reversed sequence if the exc_type is None, otherwise, return the exception type and message in a nicely formatted manner.

Exception occurred: <class 'AttributeError'>
Exception message: 'str' object has no attribute 'copy'
Execution of the program continues...

The exception was handled correctly by the __exit__ method, and because we returned True when the error occurs, the program execution continues even after exiting the with statement block and we know because the print statement was executed which is written outside the with block.

Conclusion

Context managers provide a way to manage resources efficiently like by preparing them to use and then releasing them after they are no longer needed. The context managers can be used with Python's with statement to handle the setup and teardown of resources in the program.

However, we can create our own custom context manager by implementing the enter(setup) logic and exit(teardown) logic within a Python class.

In this article, we've learned: