Context Managers in Python

Most of the Python programmers must have, either knowingly or unknowingly have dealt with a Context Manager.

Python provides such ease of using abstractions that most of the time we overlook the clever feature thinking it’s a bit trivial abstraction. Context Managers are no exception in that case. 

If we look into the Python Docs this is what we would find.

A context manager is an object that defines the runtime context to be established when executing a with statement. - Python Docs

This might be a bit much for most of us, so in simple terms:

Objects that help with control flow of the with statement - Luciano Ramalho, Fluent Python

The most common Context Manager that every Python programmer uses very frequently for managing files, is with the open function

with open('StarWars.txt') as f:
    f.write('This is the way!')
  1. Here a temp context is set when the open function returns an IO object here f

  2. You write to the file with the IO object f

  3. When all things are done, that temp context is reliably closed.  Somewhat like a setup/teardown process.

You must think that it is fine and all but why can’t I use a simple try/finally block in a similar manner. Well technically can do that too.

try:
  file = open('StarWars.txt')
  file.write('This is the way!')
except FileNotFoundError as e:
  logging.error('File not found', exec_info=True)
finally:
  file.close()

But then you would be defeating the whole purpose of a Context Manager.

Context Managers exists to control a with statement and the with statement was designed to simplify the try/finally pattern.

Besides that, a Context Manager can help in improving code readability of the code and abstracting complex logic with a single withas control flow.

The withas control flow

Now that we have some basic idea of what a Context Manager is? and why do they exist? Let's come back to our file example and take a bit of a deeper dive to see how the withas control flow is executed. 

with open('StarWars.txt') as f:
    f.write('This is the way!')
  1. With statement is called with an expression.

  2. The result of that expression is the value of the variable after as in the withas statement here f

  3. Some work is done inside the with block with f inside the with block.

  4. Finally when the control flow exits the with block the file is closed.

In other cases this closing of resource can also be releasing a resource that was being used or a previous state is being restored that was being changed inside the with block.

Writing Your Own Context Manager

Context Manager can be written in two ways either create your own class or use the Contextlib module from the standard lib.

Creating a class with magic methods:

class FileManager:
    def __init__(self, filename, mode='r'):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.open_file = Open(self.filename, self.mode)
        return self.open_file

    def __exit__(self, exception, exception_type, exception_traceback):
        self.open_file.close()

>> with FileManager('StarWars.txt', 'w') as f:
        f.write('This is the way!')

FileManager class tries to mimic the open function and exposes the two most important dunder methods __enter__ and __exit__ which are responsible for creating a temp context, doing the work, and restoring the state or object. 

  1. When the with block is called, dunder __enter__ is called with just the self as an argument. 

  2. In the enter method, an IO object i.e self.open_file is returned. This returned value is of f , that is the variable that is used after the as statement.

  3. f is used to write to the file and when the with control flow is moved out of the block and the __exit__

On a successful exit, all the three arguments of the dunder exit will have a NONE value. For some reason, if you want to suppress it, you should return True as anything returned besides True raises exception.

Now let’s look at somewhat real-world examples of Class-based Context Manager.

~ Real World Example

class DatabaseHandler:

    def __init__(self):
        self.host = 'localhost'
        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, exc_type, exc_value, traceback):
        self.cursor.close()
        self.connection.close()

DatabaseHandler is the Postgres handler that we can import to other parts of the program.

  1. __enter__ method creates a DB connection and returns a cursor object

  2. __exit__ method checks for exceptions, raises if any, and closes the connection reliably.

When using DatabaseHandler we don't have to worry much about opening/closing the connection or even managing exceptions. Which in return makes our code more readable. 

class Editor:

    def get_articles(self, start_date, end_date):
        sql = 'Select query to get total articles'
        with DatabaseHandler() as db:
            db.execute(sql, (start_date, end_date))
            total_articles = db.fetchall()
        return total_articles

    def add_journalist(self, first_name, last_name, email):
        sql = 'Insert query to add a new user to the CMS'
        with DatabaseHandler() as db:
            db.execute(sql, (first_name, last_name, email))

Using @contextmanager

The @contextmanager decorator converts your generator function into a context manager so that it can be used with the withas control flow.

from contextlib import contextmanager

@contextmanager
def foobar():
    print("Before")
    yield {}
    print("After")

with foobar() as fb:
    print(fb)

"""
---OUTPUT---
Before
{}
After
"""

Yield is used to split the function into two halves. Everything before yield will be executed at the beginning of the with block and evrerything after yield will be called at the end of the block.

Even though you can’t see the __enter__ and __exit__ explicitly here, under the hood, they are being called. 

  1. __enter__ before the yield

  2. __exit__ before the yield

To learn more about how they are implemented, I did peek into some of the module's code.

Peeking into contextlib code

Going through the source code of contextlib I did find the __enter__ method and __exit__ method implementation.

def __enter__(self):
    # do not keep args and kwds alive unnecessarily
    # they are only needed for recreation, which is not possible anymore
    del self.args, self.kwds, self.func
    try:
    	return next(self.gen)
    except StopIteration:
    	raise RuntimeError("generator didn't yield") from None

Since we are dealing with generator function here, the __enter__ method's main job is to next or yield the value of the gen object and return the yielded value so that it can be assigned to the variable after as in the withas block.

def __exit__(self, type, value, traceback):
    if type is None:
        try:
            next(self.gen)
        except StopIteration:
            return False
        else:
            raise RuntimeError("generator didn't stop")
    else:
        if value is None:
            value = type()
        try:
            self.gen.throw(type, value, traceback)
        except StopIteration as exc:
            return exc is not value
        except RuntimeError as exc:
            if exc is value:
                return False
            if type is StopIteration and exc.__cause__ is value:
                return False
            raise
        except:
            if sys.exc_info()[1] is value:
                return False
            raise
        raise RuntimeError("generator didn't stop after throw()")

The exit method checks for an exception, if anything is passed to the type besides NONE, self.gen_throw is called causing the exception to be raised in the yield line. Otherwise next(gen) is called resuming the generator function body after the yield.

To make the code a bit readable on the screen I did remove the comments from the source code, feel free to check the source code on Github. I was amazed at how well commented the code was. Kudos to the devs for keeping the codebase so well maintained.

~ Real World Example with @contexmanager

@contextmanager
def database_handler():
	    host = 'localhost'
        user = 'dev'
        password = 'dev@123'
        db = 'foobar'
        port = '5432'
        connection = None
        cursor = None
        connection = psycopg2.connect(
            user=self.user,
            password=self.password,
            host=self.host,
            port=self.port,
            database=self.db
        )
        cursor = self.connection.cursor()
        yield cursor
        cursor.close()
        connection.close()

The code snippet would work on a happy flow of the program but for any exception, this program is seriously flawed. You might remember the exit method in the contextlib source code, it raises exceptions if any in the yield line.  Here we don’t seem to have any error handling in the yield line.

So during an exception generator function will abort without closing the resource properly. Failing to do the single job that it was meant for.

Tackling exceptions when using @contextmanager

@contextmanager
def database_handler():
	try:
	    host = 'localhost'
        user = 'dev'
        password = 'dev@123'
        db = 'foobar'
        port = '5432'
        connection = None
        cursor = None
        connection = psycopg2.connect(
            user=self.user,
            password=self.password,
            host=self.host,
            port=self.port,
            database=self.db
        )
        cursor = self.connection.cursor()
        yield cursor
    except Exceptions as e:
    	logging.errror("Error connecting to the database", exc_info=True)
    finally:
        cursor.close()
        connection.close()

By default __exit__ suppresses the exception. So exceptions should be reraised in the decorator function.

Having a try/finally (or a with block) around the yield is an unavoidable price of using @contextmanager, because you never know what the users of your context manager are going to do inside their with block. - Luciano Ramalho


PyRandom

The blog post is a follow up to the talk I gave at PythonPune about Context Managers. While researching for the talk I came across some interesting things hence I am adding this to the section PyRandom.

  1. localcontext is great making high precision calculations in a particular withas block as when the flow exits the block, the precision is automatically restored.

    from decimal import localcontext
    
    with localcontext() as arthmetic:
    	# Sets the current decimal precision to 42 places
    	arthmetic.prec = 42
    	# Do some high_precision_arithmetic here
    
    # Automatically restores precision to the previous context 
    arithmetic_calculation()
  2. When using multiple context managers the control flow seems to work in a LIFO, Last In First Out manner. The __enter__ method that is called last will have it’s __exit__ method called first. 

    import contextlib
    
    @contextlib.contextmanager
    def make_context(name):
        print ('entering:', name)
        yield name
        print ('exiting :', name)
    
    with make_context('A') as A, make_context('B') as B, make_context('C') as C:
        print ('inside with statement:', A, B, C)
        
    """
    ---OUTPUT---
    entering: A
    entering: B
    entering: C
    inside with statement: A B C
    exiting : C
    exiting : B
    exiting : A
    """
    
  3. Context managers created with @contextmanager are for single use.

    from contextlib import contextmanager
    
    @contextmanager
    def foobar():
        print("Before")
        yield
        print("After")
    
    >>> foo = foobar()
    
    >>> with foo:
            pass
    
    Before
    After 
    
    >>> with foo:
            pass
    
    Traceback (most recent call last):
    RuntimeError: generator didn't yield
  4. Context Managers that can be nested with the with control flow and can use the same instance of a context manager in a nested flow are called Reetrant Context Managers.

    >>> from contextlib import redirect_stdout
    >>> from io import StringIO
    >>> stream = StringIO()
    >>> write_to_stream = redirect_stdout(stream)
    >>> with write_to_stream:
    ...     print("This is written to the stream rather than stdout")
    ...     with write_to_stream:
    ...         print("This is also written to the stream")
    ...
    >>> print("This is written directly to stdout")
    This is written directly to stdout
    >>> print(stream.getvalue())
    This is written to the stream rather than stdout
    This is also written to the stream
  5. Context Managers that cannot have a nested with control flow but their single instance can be used multiple times are called Reusable Context Managers

    from threading import Lock
    
    lock = threading.Lock()
    
    # Lock gets acquired
    with lock:
        # Do some work
    # Lock gets released
    
    # Lock gets acquired
    with lock:
        # Do some work
    # Lock gets released
    
    # Lock gets acquired
    with lock:
    	# Do some work
    	with lock:
    		# OOPS! Deadlock
    		# Try to acquire an already acquired lock 

Let's give credit where it's due. The talk and blog post both were created from the notes that I took two years back while reading Fluent Python by Luciano Ramalho.

I highly recommend that book if you want to upskill your Python Knowledge.

Happy Coding!

Write a comment ...

Pradhvan Bisht

Show your support

If you enjoy my content please consider supporting it.

Write a comment ...

Pradhvan Bisht

Yet another blog around programming, open-source software, FOSS communities and sometimes life!