Context Manager in Python

If you are new in python then you might wonder how ‘with’ keyword works. Lets decode its working today.

We all like python due to its concise and pretty syntax but internally python do some hidden work to provide this functionality. In case of context manager ‘__enter__’ and ‘__exit__’ come into picture. Yup more magic methods.

So context manager are the objects which define __enter__ and __exit__ on there objects. Most of the time context manager are use to reduce finally block code.

Let see an example.

with open('data.txt', 'w') as f:
    f.write('Hello World')

f = open('data.txt', 'w')
try:
    f.write('Hello World')
finally:
    f.close()

Code sample #1 and #2 are exactly same when come to operations but #1 looks more concise.

Note: Above code opens a file ‘data.txt’ in truncate write mode and write ‘Hello World’ into file and close that file.

We can use context managers for the management of resources like files, connections, database transaction and other.

If you remember we create a time logger decorator in our decorator tutorial. Lets create similar with context managers.

with CONTEXT_MANAGER as RESOURCE:
    # Code block that uses the RESOURCE

import time


class LogTime:
    def __init__(self, name=None):
        self.name = name

    def __enter__(self):
        self.start = time.time()
    
    def __exit__(self, *args):
        self.end = time.time()
        self.delta = self.end - self.start
        if self.name:
            print(f"{self.name} took {self.delta:.6f}s")
        else:
            print(f"took {self.delta:.6f}s")

with LogTime(name='sleep1'):
    time.sleep(1)

with LogTime():
    time.sleep(2)

sleep1 took 1.000267s
took 2.000144s

So when interpreter found with then it call __enter__ on the context_manager object and when code under with completes it call __exit__.

It looks more easy then decorators.

Lets create our own FileHandler as another example.

class FileHandler:
    def __init__(self, filename, fmode):
        self.filename = filename
        self.fmode = fmode

    def __enter__(self):
        self.fp = open(self.filename, self.fmode)
        return self.fp

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.fp.close()

with FileHandler('data.txt', 'w') as f:
    f.write('Hello World')

$ cat data.txt

Hello World

Above code is easy to understand Lets dive more deeper.

Hope you noticed exc_type, exc_val and exc_tb.

These params are automatically passed by with statement and contains information about the exception. If no exception raised these values will be None.

exc_type: Type of exception Raised

exc_val: Value passed when Exception raised.

exc_tb: traceback object to get more information regarding Exception.

 class ErrorLogger:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            # Log the exception
            print(f"Exception Type ->  '{exc_type.__name__}'\n"
                  f"value: {exc_val}")
            print(f"Traceback:\n{traceback.print_tb(exc_tb)}")
        else:
            print("No exception occurred.")

with ErrorLogger():
    raise ValueError("An error occurred!")

Exception Type ->  'ValueError'
value: An error occurred!
Traceback: ...

with ErrorLogger():
    pass

No exception occurred.

Useful use case of context manager.

from django.db import transaction

with transaction.atomic():
    # do something 
    pass 

Above written code snippet is django atomic transaction application. So when an exception is raised exit will roll the transaction.

Note: If you use try except under inside with in transaction.atmoic() in that case no rollback will happen.

Using contextlib instead of magic methods.

Although magic methods are clean but we can also create context manager without defining magic methods.

from contextlib import contextmanager

@contextmanager
def FileHandler(file_name, file_mode):
    fp = None
    try:
        fp = open(file_name, file_mode)
        yield fp
    finally:
        if fp:
            fp.close()

with FileHandler('data.txt', 'w') as f:
    f.write('Hello World')

This code is replacement of above FileHandler but function based implementation. Lets understand working.

1. When first with call FileHandler then yield will return file_object and pause the execution of FileHandler.
2. And when statements under with got executed then this yield will act as entry point for __exit__ and perform the operation.

Async Context Manager

Consider that i want to use context manager with asynchronous code written with async and await.

import asyncio

async with LogTime(name='sleep1'):
    await asyncio.sleep(1)

TypeError: 'LogTime' object does not support the asynchronous context manager protocol

Oh we have a problem so python provide two other magic methods __aenter__ and __aexit__ for handling async code in context manager.

import time
import asyncio

class LogTime:
    def __init__(self, name=None):
        self.name = name

    async def __aenter__(self):
        self.start = time.time()

    async def __aexit__(self, *args):
        self.end = time.time()
        self.delta = self.end - self.start
        if self.name:
            print(f"{self.name} took {self.delta:.6f}s")
        else:
            print(f"took {self.delta:.6f}s")


async with LogTime(name='sleep1'):
    await asyncio.sleep(1)

sleep1 took 1.001438s

So we fixed it with async method.

Leave a clap if you found this helpful 👏 😃.

In Plain English

Thank you for being a part of our community! Before you go:

• Be sure to clap and follow the writer! 👏
• You can find even more content at PlainEnglish.io 🚀
• Sign up for our free weekly newsletter. 🗞️
• Follow us on Twitter(X), LinkedIn, YouTube, and Discord.