Python: Guide to Context Manager & How do they work?
Tags: python, programming
Published: 2024-02-27 14:20:46

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.

```python
# 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.

```python
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.

```python
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.

```python
# 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**

```python
# 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.

```plaintext
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.

```python
# 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.

```bash
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.

```python
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.

```bash
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?**

```python
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.

```bash
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.

```python
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.

```bash
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:

* **What is context manager and why they are used**
    
* **Using context manager with the** `with` **statement**
    
* **Implementing context management protocol within a class**
    

---

🏆**Other articles you might be interested in if you liked this one**

✅[Understanding the basics of abstract base class(ABC) in Python](https://geekpython.in/abc-in-python).

✅[Implement \_\_getitem\_\_, \_\_setitem\_\_ and \_\_delitem\_\_ in Python class to get, set and delete items](https://geekpython.in/implement-getitem-setitem-and-delitem-in-python).

✅[Generate and manipulate temporary files using tempfile in Python](https://geekpython.in/tempfile-in-python).

✅[Using match-case statement for pattern matching in Python](https://geekpython.in/match-case-in-python).

✅[Comparing the sort() and sorted() function in Python](https://geekpython.in/python-sort-vs-sorted).

✅[Using super() function to implement attributes and methods of the parent class within the child class](https://geekpython.in/super-in-python).

✅[Using str and repr to change string representation of the objects in Python](https://geekpython.in/str-and-repr-in-python).

---

**That's all for now**

**KeepCoding✌✌**