Writing Custom Context Managers in Python

Motivation

You might already be familiar with the with statement, a concise way to handle resources in Python. But have you ever wondered how it works under the hood? Well, the power of the with keyword comes from context managers.

Context managers are a fundamental design pattern in Python that provide a structured approach to resource management. They ensure that resources are acquired, used properly, and then finally released or cleaned up, even with errors or exceptions.

This becomes especially important when dealing with resources like files, network connections, or database handles. By using context managers, we can write cleaner and more reliable code, freeing ourselves from the worry of forgetting to close a file or release a lock.

In this tutorial, we will go beyond the default provided context managers of Python and learn to write our own.

Understanding Context Managers in Python

Under the hood, context managers are objects that define two special methods: __enter__ and __exit__. The __enter__ method is called when you enter the with block, and its return value is assigned to a variable within that block. The __exit__ method, on the other hand, is called when the with block exits, regardless of whether it finishes normally or with an exception.

This structure ensures proper resource handling. Let’s look at some built-in context managers in Python to understand this more:

1. File Management

Take the classic example of opening a file:

with open('file.txt', 'r') as file:
   data = file.read()

Here, open('file.txt', 'r') acts as the context manager. When you enter the with block, the enter method of the file object is called, opening the file and assigning it to the variable file. You can then use file.read() to access the file contents.

Next, the __exit__ method of the file object is guaranteed to be called when the with block exits, even if an exception occurs. This method takes care of closing the file, ensuring you don't leave open file handles lying around.

2. Thread Locking

Moving beyond files, context managers can also be used for thread synchronization using threading.Lock():

import threading

lock = threading.Lock()
with lock:
   # Critical section
   print("This code is executed under lock protection.")

Here, lock is a threading.Lock object, another context manager. When you enter the with block, the __enter__ method acquires the lock, ensuring only one thread can execute the critical section at a time. Finally, the __exit__ method releases the lock upon exiting the with block, allowing other threads to continue.

3. Database Connections

Similarly, context managers can manage database connections:

import sqlite3

with sqlite3.connect('database.db') as connection:
    cursor = connection.cursor()
    cursor.execute("SELECT * FROM table")
    rows = cursor.fetchall()

The sqlite3.connect('database.db') call is a context manager. Entering the with block establishes a connection to the database, assigning it to connection. You can then use a cursor to interact with the database. The __exit__ method guarantees the connection is closed when the with block exits, preventing resource leaks.

4. Network Sockets

Context managers can even handle network communication:

import socket

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    s.connect(('localhost', 8080))
    s.sendall(b'Hello, world')
    data = s.recv(1024)

Here, socket.socket(socket.AF_INET, socket.SOCK_STREAM) creates a socket object that acts as the context manager. The __enter__ method creates the socket, and within the with block, you can connect, send data, and receive data. The __exit__ method ensures the socket is closed properly when done.

5. Directory Scanning

This time, os.scandir('.') provides a way to iterate over directory entries:

with os.scandir('.') as entries:
   for entry in entries:
       print(entry.name)

os.scandir('.') acts as the context manager here. The __enter__ method opens a directory scan, and you can iterate over the entries within the with block. The __exit__ method cleans up the directory scan upon exiting.

As you can see, the with statement, powered by context managers, simplifies resource management by handling allocation and deallocation automatically.

Writing custom context managers in Python

Now, let’s finally look at how to write your own custom managers to have fine-grained control over resource management.

There are two main approaches to writing custom context managers: class-based and function-based.

Class-Based Approach

The class-based approach is the most structured and flexible method for writing context managers. Here, you define a class that implements the special methods __enter__ and __exit__.

Let’s look at an example of a Timer class that measures execution time:

import time


class Timer:
   def __enter__(self):
       self.start_time = time.time()
       return self

   def __exit__(self, exc_type, exc_value, traceback):
       self.end_time = time.time()
       elapsed_time = self.end_time - self.start_time
       print(f"Elapsed time: {elapsed_time} seconds")


# Example usage
if __name__ == "__main__":
   with Timer() as timer:
       # Code block to measure the execution time
       time.sleep(2)  # Simulate some time-consuming operation
Elapsed time: 2.002082347869873 seconds

Here, the Timer class defines the __enter__ method to capture the start time when you enter the with block. It returns self to allow access to the object within the block. The __exit__ method calculates the elapsed time upon exiting the with block and prints it.

Function-Based Approach

If you prefer a more concise approach possibly at the cost of flexibility, the function-based approach might be a good option. Here, you use the contextmanager decorator from the contextlib module to convert any function into a context manager.

Let’s see how that works:

import time
from contextlib import contextmanager


@contextmanager
def timer():
   start_time = time.time()
   yield
   end_time = time.time()
   elapsed_time = end_time - start_time
   print(f"Elapsed time: {elapsed_time} seconds")


# Example usage
if __name__ == "__main__":
   with timer():
       time.sleep(2)
Elapsed time: 2.0020740032196045 seconds

The @contextmanager decorator transforms the timer function into a context manager. Inside the function, start_time is captured, and the yield statement pauses execution, allowing code within the with block to run.

Finally, __exit__ functionality is achieved by capturing the end time and printing the elapsed time.

Essentially, you write the logic for the __enter__ before the yield keyword whereas the logic for __exit__ comes after. Both approaches achieve the same outcome, but the choice depends on your preference for structure and readability.

Practical Examples of Context Managers

Now that you understand how to write custom context managers in Python, let’s look at some practical examples where they can be incredibly useful.

Managing File Operations

File operations are a common task in many applications. Whether reading from or writing to files, proper management ensures data integrity and resource efficiency. Let’s consider a scenario where we want to write text to a file using a custom context manager.

class FileManager:
   def __init__(self, filename, mode):
       self.filename = filename
       self.mode = mode
       self.file = None

   def __enter__(self):
       self.file = open(self.filename, self.mode)
       return self.file

   def __exit__(self, exc_type, exc_value, traceback):
       if self.file:
           self.file.close()


# Example usage
if __name__ == "__main__":
   with FileManager("example.txt", "w") as file:
       file.write("Hello, world!\n")

In this example, we define a FileManager class to mimic the open function in Python that handles file operations. Let's break down what each method does:

• __init__(self, filename, mode): This method is the constructor of the class. It initializes the FileManager object with the provided filename and mode, which specify the file to be opened and the mode in which it should be opened (e.g., read, write, append).
• __enter__(self): This method is called when entering the with block as always. It opens the file specified by filename in the mode specified by mode. The opened file object is assigned to self.file and returned to be used within the with block.
• __exit__(self, exc_type, exc_value, traceback): This method is called when exiting the with block, regardless of whether an exception occurred within the block. It ensures that the file is closed properly by calling self.file.close() if self.file is not None. This prevents resource leaks and ensures proper cleanup.

If you inspect the contents of example.txt after executing the above code, you should see the written text:

$ cat example.txt
Hello, world!

Of course, the advantage of FileManager class is that you can implement any custom logic to work with files. For example, you can modify it so that the class keeps a log of which files are opened or closed to a text file.

import datetime


class FileManager:
   def __init__(self, filename, mode, log_filename="file_log.txt"):
       self.filename = filename
       self.mode = mode
       self.log_filename = log_filename
       self.file = None

   def log_action(self, action):
       timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
       with open(self.log_filename, "a") as log_file:
           log_file.write(f"{timestamp} - {action}: {self.filename}\n")

   def __enter__(self):
       self.file = open(self.filename, self.mode)
       self.log_action("Opened")
       return self.file

   def __exit__(self, exc_type, exc_value, traceback):
       if self.file:
           self.file.close()
           self.log_action("Closed")


# Example Usage
with FileManager("example.txt", "r") as file:
   content = file.read()
   print(content)

Hello, world!

In this modified version:

• datetime module is imported to work with timestamps.
• The log_action method is added to log file actions.
• log_filename parameter is added to specify the filename for the log file. The default is set to 'file_log.txt'.
• The filename and action (opened or closed) are appended to the log file using the log_action method within __enter__ and __exit__ methods.
• Inside the log_action method, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") is used to get the current timestamp in the format of "YYYY-MM-DD HH:MM:SS".
• The timestamp is then included in the log along with the filename and action.

Now, if we print the contents of file_log.txt, we should see the logs:

$ cat file_log.txt
2024-04-03 12:37:02 - Opened: example.txt
2024-04-03 12:37:02 - Closed: example.txt

Another practical use case we can implement as a custom context manager is an SQLite database management tool:

import sqlite3


class DatabaseConnection:
   def __init__(self, database_name):
       self.database_name = database_name
       self.connection = None

   def __enter__(self):
       self.connection = sqlite3.connect(self.database_name)
       return self.connection

   def __exit__(self, exc_type, exc_value, traceback):
       if self.connection:
           self.connection.close()

Again, this mimics the sqlite3.connect function, but it can be further improved by adding custom database management logic. Let's check if the manager is working with some sample functions:

# Example usage
def create_table():
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute(
           """CREATE TABLE IF NOT EXISTS users (
                           id INTEGER PRIMARY KEY,
                           username TEXT,
                           email TEXT)"""
       )


def insert_data(username, email):
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute(
           "INSERT INTO users (username, email) VALUES (?, ?)", (username, email)
       )
       connection.commit()


def fetch_data():
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute("SELECT * FROM users")
       return cursor.fetchall()


if __name__ == "__main__":
   create_table()
   insert_data("john_doe", "john@example.com")
   insert_data("jane_doe", "jane@example.com")
   users = fetch_data()
   print("Users in the database:")
   for user in users:
       print(user)

Users in the database:
(1, 'john_doe', 'john@example.com')
(2, 'jane_doe', 'jane@example.com')
(3, 'john_doe', 'john@example.com')
(4, 'jane_doe', 'jane@example.com')

We can write much more maintainable code by breaking up the database operations into functions and using the custom manager we just created.

Advanced Topics Related to Context Managers in Python

In this section, we’ll discuss advanced topics related to context managers in Python. These include error handling and nesting context managers.

Error handling

Error handling is an essential aspect of programming, and context managers in Python provide a convenient way to manage errors within resource management operations.

In the FileManager class example, error handling can be implemented within the __exit__ method to ensure proper cleanup even if an exception occurs. For instance, if an error occurs while writing data to the file (while inside the with block), we still want to ensure the file is closed to prevent resource leaks. We can modify the __exit__ method as follows:

def __exit__(self, exc_type, exc_value, traceback):
   if self.file:
       self.file.close()
       if exc_type is not None:
           print(f"An error occurred: {exc_value}")

Here, we first close the file using self.file.close() to ensure proper cleanup. Then, we check if an exception occurred (exc_type is not None). If so, we print an error message to provide feedback on what went wrong.

Here:

• exc_type is the exception class, such as AttributeError and IndexError.
• exc_value is the exception instance with the message.
• traceback is the full traceback of the error.

Let’s run the FileManager class with a deliberately incorrect code:

with FileManager("example.txt", "r") as file:
   # Simulate a sample error
   content = file.read()
   print(content[150])

An error occurred: string index out of range

So, to implement custom logic when handling exceptions in your custom managers, here is a sample workflow you can follow:

def __exit__(self, exc_type, exc_value, traceback):
   # Perform cleanup
   ...

   # Handle the exception
   if isinstance(exc_value, AttributeError):
       # Handle AttributeError here
       print(f"Exception message: {exc_value}")
   # Handle another exception
   elif isinstance(exc_value, AnotherError):
       ...

Nesting context managers

Nesting context managers allow you to manage multiple resources simultaneously, each with its own context manager, within a single with statement. This is particularly useful when dealing with complex operations that involve multiple resources. Let's see how we can nest context managers in our examples.

Consider a scenario where we want to open a file and perform database operations within the same context. Perhaps, we want to input the contents of a CSV file as rows into a database. We can achieve this by nesting the FileManager and DatabaseConnection context managers within a single with statement:

if __name__ == "__main__":
   with FileManager("example.txt", "w") as file, DatabaseConnection("example.db") as connection:
       ...

In this example, the FileManager and DatabaseConnection context managers are nested within the same with statement. This allows us to open the file and establish a database connection simultaneously. Subsequent file writing and database operations are performed within this combined context.

One thing to keep in mind is the readability of code. I would recommend that the number of simultaneous contexts not exceed two.

Conclusion

Context managers in Python offer a structured approach to resource management within a with statement, ensuring proper allocation and deallocation of resources.

Throughout this tutorial, we’ve learned the basics of context managers, went through practical examples, and discussed advanced topics like error handling and nesting. By mastering writing custom context managers, you can write cleaner, more reliable code, improve error handling, and manage resources efficiently.

The topic of writing custom context managers falls under the umbrella of “Advanced Python” category. Here are some related resources to further enhance your skills:

• Writing functions in Python course
• Becoming a Python Programmer skill track