Manav Garg

Manav Garg

Context Managers in Python.

Context Managers in Python.

Subscribe to my newsletter and never miss my upcoming articles

This post is inspired from a talk at python pune january meetup by Pradhvan.

What are Context Managers?

Here's what Python's official documentation says:

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

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:

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:

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().

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()

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.

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

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:

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

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

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()

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:

 
Share this