Context Managers in Python.
Tags: python, programming, beginners
Published: 2021-01-29 12:25:21

This post is inspired from a talk at [python pune january meetup](https://www.youtube.com/watch?v=cEYwfeAXGwQ) by [Pradhavan](https://pradhvanbisht.in/). 

# What are Context Managers?
Here's what Python's official documentation says:
*<blockquote>A context manager is an object that defines the runtime context to be established when executing a with statement. The context manager handles the entry into, and the exit from, the desired runtime context for the execution of the block of code. Context managers are normally invoked using the with statement, but can also be used by directly invoking their methods. – [Python Docs](https://docs.python.org/3/reference/datamodel.html#with-statement-context-managers)</blockquote>*

But that's just covering all the bases. Let's understand a simplified version. 

Programmers work with external resources from time to time, like files, database connections, locks etc. Context managers allow us to manage those resources by specifying:
1. What to do when we acquire the resource, and
2. What to do when the resource gets released

# Why do we need Context Managers?
Consider the following example:
```python
for _ in range(100000):
    file = open("foo.txt", "w")
    files.append(f)
    file.close()
```
Notice that we're calling the `close()` method to ensure that the file descriptor is released every time. If we didn't do that, our OS would run out of its allowed limit to open file descriptors eventually.

However, we write a more pythonic version of the above code using context manager:
```python
for _ in range(100000):
    with open("foo.txt", "r") as f:
        files.append(f)
```
Here `open("foo.txt", "r")` is the context manager that gets activated using the `with` statement. **Notice that we didn't need to explicitly close the file, the context manager took care of it for us.** Similarly there are other predefined context managers in Python that makes our work easier.

# Can we define our own Context Manager?
Yes. There are two ways to define a custom context manager:
1. Class based definition.
2. Function based definition.

# Class Based Context Managers
Let's continue with our file example and try to define our own context manager which will emulate `open()`.

```python
files = []

class Open():
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.open_file = open(self.filename, self.mode)
        print("Enter called")
        return self.open_file

    def __exit__(self, *args):
        print("Exit Called")
        self.open_file.close()

for _ in range(100000):
    with Open('foo.txt', 'w') as f:
        files.append(f)
```
* The `__enter__` method tells us what to do when we acquire the resource, i.e., giving us a file object.
* The `__exit__` method defines what to do when we're exiting the context manager, i.e., closing the file.
* You can see how both `__enter__` and `__exit__` are called with every loop. 

### Handling Errors
**How do we handle `FileNotFoundError` with python's `open()`**
```python
try:
    with open("foo.txt", "r") as f:
        content = f.readlines()
except FileNotFoundError as e:
    print("Hey, file isn't there. Let's log it.")
```
Such a basic error handling code that needs to be every time you open a file. Let's try to DRY it with our custom context manager.
```python
class Open():
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        print("Enter called")

        try:
            self.open_file = open(self.filename, self.mode)
            return self.open_file
        except FileNotFoundError as e:
            print("Hey, file isn't there. Let's log it.")

    def __exit__(self, exc_type, exc_value, exc_traceback): #notice the parameters
        print("Exit Called")
        if(exc_type is None):
            self.open_file.close()
            return True
        else:
            return True

with Open("foo.txt", "r") as f:
    content = f.readlines()
```
**Changes in `__exit__`**
* `exc_type` is type of error Class which you'll get while handling errors in `__enter__` (AttributeError in this case).
* `exc_value` is the value of the error which you'll get while handling errors in `__enter__`.
* `exc_traceback` is the traceback of the error which you'll get while handling errors in `__enter__`.
* We're returning `True` to suppress the error traceback (not to be confused with `exc_traceback` parameter).

### Another Real World Example
```python
class DatabaseHandler():
    def __init__(self):
        self.host = '127.0.0.1'
        self.user = 'dev'
        self.password = 'dev@123'
        self.db = 'foobar'
        self.port = '5432'
        self.connection = None
        self.cursor = None

    def __enter__(self):
        self.connection = psycopg2.connect(
            user=self.user,
            password=self.password,
            host=self.host,
            port=self.port,
            database=self.db
        )
        self.cursor = self.connection.cursor()
        return self.cursor

    def __exit__(self, *args):
        self.cursor.close()
        self.connection.close()
```
# Function Based Context Managers
Function based context management is done by using a lib called `contextlib`, through which we can change a simple generator function into a context manager. Here's what a typical blueprint looks like:
```python
from contextlib import contextmanager

@contextmanager
def foobar():
    print("What you would typically put in __enter__")
    yield {}
    print("What you would typically put in __exit__")

with foobar() as f:
    print(f)
```
* `contextmanager` decorator is used to turn any generator function into a context manager.
* `yield` work as a separater between `__enter__` and `__exit__` parts of the context manager.  

**Handling files**
```python
from contextlib import contextmanager

@contextmanager
def open_(filename, mode):
    
    print("SETUP")
    open_file = open(filename, mode)

    try:
        print("EXECUTION")
        yield open_file

    except:
        print("Hey, file isn't there. Let's log it.")

    finally:
        print("CLEAN-UP")
        open_file.close()


with open_("somethign.txt", "w") as f: #notice the mode
    content = f.readlines() #you cannot read on write mode
```
We wrap `yield` in a `try` block because we don't know what the user is going to do with the file object. They might try to use it in a way that it's not intended to (as shown above).

**Database Connections**
```python
from contextlib import contextmanager

@contextmanager
def database_handler():
    try:
        host = '127.0.0.1'
        user = 'dev'
        password = 'dev@123'
        db = 'foobar'
        port = '5432'
        connection = psycopg2.connect(
            user=user,
            password=password,
            host=host,
            port=port,
            database=db
        )
        cursor = connection.cursor()
        yield cursor
    
    except:
        print("Hey, file isn't there. Let's log it.")

    finally:
        cursor.close()
        connection.close()
```
<hr>

# Resources
We have just covered just an introduction to context managers, but I feel that it's just the tip of the iceberg and there are many interesting use cases for it. Here are some interesting links that I found:
* https://docs.python.org/3/reference/datamodel.html#context-managers
* https://docs.python.org/3/library/contextlib.html
* https://rednafi.github.io/digressions/python/2020/03/26/python-contextmanager.html
* https://www.markus-gattol.name/ws/python.html#context_manager