Python context managers are powerful

A context manager is an object that defines the methods __enter__() and __exit__() to set up and tear down a context for a block of code. It allows for the automatic and consistent management of resources, such as file handles or network connections.

For example, when a customer enters a restaurant, they are greeted by the host and seated at a table. This is similar to the __enter__() method of a context manager, where resources are acquired and set up for the customer to use.

The customer then orders their meal and the kitchen staff prepare it for them. This is similar to the block of code that is executed within the context of the manager.

When the meal is ready, it is served to the customer on a plate. This is similar to the resources that are made available to the block of code.

After the customer has finished their meal, they ask for the bill and pay it before leaving the restaurant. This is similar to the __exit__() method of a context manager, where the resources are cleaned up and released.

Just as the restaurant staff takes care of seating the customer, serving the meal, and cleaning up after the customer leaves, a context manager takes care of acquiring resources, making them available to the block of code, and cleaning up and releasing the resources when the block of code has finished executing.

In a restaurant, even if the customer leaves the restaurant without paying the bill, the staff will still clean up and prepare the table for the next customer, similarly in a context manager, regardless of whether the block of code completes successfully or raises an exception, the resources are always cleaned up and released by the __exit__() method.

Context managers were not part of the original implementations in python, but as we always want an efficient way to deal with programming tasks, the context manager where added. In Golang, Dart and few other languages, context managers where extended to pass data across executions, now let’s look at it’s history with python.

Context managers were first introduced in Python 2.5, in 2005. They were added as a way to simplify the management of resources, such as file handles and network connections, that need to be acquired and released in a consistent and reliable way.

Prior to the introduction of context managers, acquiring and releasing resources required a lot of boilerplate code and was prone to errors. For example, when working with a file, a programmer would have to manually open the file, perform the necessary operations, and then close the file, remembering to include error-handling code in case an exception was raised.

Context managers provided a way to automate this process and make it more reliable by using the with statement. The with statement creates a context in which the __enter__() method of the context manager is called at the beginning of the block, and the __exit__() method is called at the end of the block, regardless of whether the block of code completed successfully or raised an exception.

This made it much easier for programmers to manage resources and ensure that they were always released, even if an exception was raised. The with statement and context managers quickly became a widely used feature of the Python programming language and are now used in many different scenarios such as file handling, threading, and even database connections.

In python 3.3, the contextlib module was added which provides some useful context manager functions such as closing, suppress and ExitStack which allow more complex context management.

Here are some of the benefits of using context:

1. Reliability: Context managers ensure that resources are always acquired and released in a consistent and reliable way, reducing the risk of errors and memory leaks.
2. Simplicity: Context managers provide a simple and readable way to manage resources, eliminating the need for boilerplate code and making it easy to understand the intent of the code.
3. Error Handling: Context managers automatically handle errors and exceptions, releasing resources even if an exception is raised, which helps to prevent resource leaks.
4. Nesting: Context managers can be nested, which allows for more complex resource management scenarios.
5. Redirecting: The contextlib package provides the contextlib.redirect_stdout and contextlib.redirect_stderr functions which can be used to redirect the standard output and standard error.
6. Mocking: The contextlib package provides the contextlib.redirect_stdin function which can be used to redirect the standard input, allowing you to mock the user input.
7. Testing: The contextlib package provides the contextlib.redirect_stdout and contextlib.redirect_stderr functions which can be used to redirect the standard output and standard error to a buffer, allowing you to capture the output for testing or debugging purposes.

Let’s Look A Custom Implementation

class MyContextManager:
    #initialize the class with the resource
    def __init__(self, resource):
        self.resource = resource

    #method that is called when the with statement is called
    def __enter__(self):
        print("Acquiring resource")
        #returns the resource so it can be used in the block
        return self.resource

    #method that is called when the block of code execution is finished
    def __exit__(self, exc_type, exc_value, traceback):
        print("Releasing resource")
        self.resource.release()

class MyResource:
    def __init__(self):
        self.acquired = False

    def acquire(self):
        print("Resource acquired")
        self.acquired = True

    def release(self):
        if self.acquired:
            print("Resource released")
            self.acquired = False

#custom function that mimics the with statement
def my_with(context_manager, block):
    #call the __enter__ method to acquire the resource
    resource = context_manager.__enter__()
    try:
        #call the block of code with the acquired resource
        block(resource)
    finally:
        #call the __exit__ method to release the resource
        context_manager.__exit__(*sys.exc_info())

# Usage
resource = MyResource()
with MyContextManager(resource) as r:
    r.acquire()
#or
my_with(MyContextManager(resource),lambda r: r.acquire())

Here, we have defined a custom context manager class MyContextManager, which takes a resource as an argument and defines the __enter__() and __exit__() methods. The __enter__() method is responsible for acquiring the resource and returning it, while the __exit__() method is responsible for releasing the resource.

We also defined a custom resource class MyResource which has acquire and release methods.

We also defined my_with function which takes the context manager and the block of code to execute as the arguments. This function is similar to the with statement, it calls the __enter__() method of the context manager to acquire the resource, then it calls the block of code with the acquired resource as the argument and finally call the __exit__() method of the context manager to release the resource.

You can use the my_with function in similar way as you use the with statement, which simplifies the process of acquiring and releasing resources and makes the code more readable.

Let’s look at few examples where we can use the with statement.

1. File handling:

In the example below, the open() function returns a file object, which acts as the context manager. The __enter__() method of the file object is called when the with statement is executed, which opens the file, and returns the file object. The block of code reads the data from the file and performs operations on it. The __exit__() method of the file object is called when the block of code is finished executing, which closes the file handle.

with open("file.txt", "r") as file:
    data = file.read()
    # perform operations on data
    # ...

2. Database connections:

import sqlite3

with sqlite3.connect("mydatabase.db") as conn:
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM users")
    result = cursor.fetchall()
    # perform operations on result
    # ...

Here, the sqlite3.connect() function returns a connection object, which acts as the context manager. The __enter__() method of the connection object is called when the with statement is executed, which establishes a connection to the database, and returns the connection object. The block of code performs a SELECT query, fetches the results and perform operations on it. The __exit__() method of the connection object is called when the block of code is finished executing, which closes the connection to the database.

3. Locking resources:

import threading

lock = threading.Lock()

with lock:
    # critical section of code
    # ...

Here also, the threading.Lock() returns a lock object, which acts as the context manager. The __enter__() method of the lock object is called when the with statement is executed, which acquires the lock, and returns the lock object. The block of code is executed only if the lock is acquired successfully, this ensures that the block of code is executed by only one thread at a time. The __exit__() method of the lock object is called when the block of code is finished executing, which releases the lock.

4. Timeout:

import signal

class TimeoutError(Exception):
    pass

def timeout_handler(signum, frame):
    raise TimeoutError("Timeout")

with signal.signal(signal.SIGALRM, timeout_handler):
    signal.alarm(5)
    try:
        # code that should be interrupted after 5 seconds
    except TimeoutError:
        pass
    finally:
        signal.alarm(0)

we use the signal module to set up a timeout. The signal.signal(signal.SIGALRM, timeout_handler) sets up a handler function that will be called when the alarm signal is raised. The signal.alarm(5) sets the alarm to go off after 5 seconds. The block of code that should be interrupted after 5 seconds is placed in the try block. The timeout_handler function raises a TimeoutError when the alarm signal is raised. Finally, the alarm is turned off by calling signal.alarm .

Let’s handle context with contextlib built-in package:

contextlib is a built-in Python package that provides a set of utility functions for working with context managers. It was introduced in Python 3.1 and provides a number of useful functions that can be used to create, modify, and compose context managers.

Let’s look at some of it’s functions:

1. contextlib.contextmanager: This is a decorator that can be used to define a context manager using a generator function. The generator should yield exactly once and the value returned is passed to the __enter__ method.

import contextlib

@contextlib.contextmanager
def my_context_manager():
    print("Entering context")
    yield "Resource"
    print("Exiting context")

with my_context_manager() as resource:
    print("Using resource:", resource)

2. contextlib.closing: This function takes an object that has a close method and returns a context manager that automatically closes the object when the context is exited.

import urllib.request

with contextlib.closing(urllib.request.urlopen("https://doziesiky.com")) as page:
    print(page.read())

3. contextlib.suppress: This function takes one or more exception classes, and creates a context manager that suppresses the specified exceptions when they are raised within the context.

with contextlib.suppress(FileNotFoundError):
    with open("file.txt") as file:
        print(file.read())

4. contextlib.ExitStack: This class provides a way to manage a stack of context managers. It can be used to create nested contexts and to exit multiple contexts at once.

with contextlib.ExitStack() as stack:
    file1 = stack.enter_context(open("file1.txt"))
    file2 = stack.enter_context(open("file2.txt"))
    # perform operations on file1 and file2
    # ...

These functions provided by the contextlib package can help simplify the process of creating, modifying, and composing context managers, making it easier to write clean, readable, and reliable code that manages resources.

Let’s redirect our outputs using context

One advanced use of the contextlib package is the ability to create a context manager that provides a temporary redirect of standard output or standard error. The contextlib package provides the contextlib.redirect_stdout and contextlib.redirect_stderr functions which can be used for this purpose.

import sys
from contextlib import redirect_stdout, redirect_stderr

with open("output.txt", "w") as f, redirect_stdout(f), redirect_stderr(f):
    print("This will be written to output.txt")
    print("And this too", file=sys.stderr)

Here, the redirect_stdout and redirect_stderr functions are used to redirect the standard output and standard error to a file. Inside the with block, any print statements will be written to the file, rather than the console. This can be useful for capturing the output of a script or library for later analysis, testing, or debugging.

Another example is redirecting the output to a string buffer, which can be useful for testing or capturing the output of a function call for further analysis.

from contextlib import redirect_stdout
import io

def my_function():
    print("Hello World!")

f = io.StringIO()
with redirect_stdout(f):
    my_function()
output = f.getvalue()
print(output)

the redirect_stdout function is used to redirect the standard output to a string buffer, which is created using the io.StringIO() function. Inside the with block, the my_function is called which has a print statement. The output of the print statement is written to the buffer instead of the console, so it can be captured and stored in the variable output for further analysis.

You can also redirect the output to the stdin using the redirect_stdin function, to simulate the user input for example:

from contextlib import redirect_stdin
import io

def my_input_function():
    return input("Enter something:")

f = io.StringIO("my input")
with redirect_stdin(f):
    result = my_input_function()
print(result)

Overall, context managers are an essential tool for Python programmers that can be used to improve the reliability, simplicity, and maintainability of code that manages resources.