How to create custom context managers in Python

Context managers in Python allow you to work more efficiently with resources, making it easier to set up and tear down resources even if there are errors when working with them. In the tutorial on how to write efficient Python code, I covered what context managers are and why they are useful. And in 3 Interesting Uses of Python Context Managers, I go over using context managers to manage threads, database connections, and more.

In this tutorial, you will learn how to create your own custom context managers. We'll review how context management works and then look at the different ways you can write your own. Let us begin.

What are context managers in Python?

Context managers in Python are objects that allow the management of resources such as file operations, database connections, or network sockets within a controlled block of code. They ensure that resources are correctly initialized before the code block is executed and automatically cleaned up afterward, regardless of whether the code block completes normally or raises an exception.

In general, context managers in Python have the following two special methods: __enter__() and __exit__(). These methods define the behavior of the context manager when entering and exiting a context.

How do context managers work?

When working with resources in Python, you should consider configuring the resource, anticipating errors, implementing exception handling, and finally releasing the resource. To do this, you will probably use a try-except-finally block like this:

try: 
    # Setting up the resource
    # Working with the resource
except ErrorType:
    # Handle exceptions
finally:
    # Free up the resource

Basically, we try to provision and work with the resource, except for any errors that may arise during the process, and finally release the resource. He finally The block is always executed regardless of whether the operation is successful or not. But with context managers and the with declaration, can have reusable try-except-finally blocks.

Now let's review how context managers work.

Enter phase (__enter__() method):
When a with a statement is found, the __enter__() The context manager method is invoked. This method is responsible for initializing and configuring the resource, such as opening a file, establishing a connection to the database, etc. The value returned by __enter__() (if applicable) is available for the context block after the “as” keyword.

Run the code block:
Once the resource is configured (after __enter__() is executed), the block of code associated with the with the statement is executed. This is the operation you want to perform on the resource.

Exit phase (__exit__() method):
After the code block completes execution, either normally or due to an exception, the __exit__() The context manager method is called. He __exit__() The method handles cleanup tasks, such as closing resources. If an exception occurs within the code block, information about the exception (type, value, traceback) is passed to __exit__() for error handling.

To sum up:

• Context managers provide a way to manage resources efficiently by ensuring that resources are properly initialized and cleaned up.
• We use the with declaration to define a context where resources are managed.
• He __enter__() The method initializes the resource and the __exit__() The method cleans up the resource once the context block is completed.

Now that we know how context managers work, let's proceed to writing a custom context manager to handle database connections.

Create custom context managers in Python

You can write your own context managers in Python using one of the following two methods:

1. Write a class with __enter__() and __exit__() methods.
2. Using the contextlib module that provides the contextmanager decorator to write a context manager using builder functions.

1. Write a class with the __enter__() and __exit__() methods

You can define a class that implements the two special methods: __enter__() and __exit__() which control the installation and removal of resources respectively. Here we write a ConnectionManager class that establishes a connection to an SQLite database and closes the connection to the database:

import sqlite3
from typing import Optional

# Writing a context manager class
class ConnectionManager:
    def __init__(self, db_name: str):
        self.db_name = db_name
        self.conn: Optional(sqlite3.Connection) = None

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_name)
        return self.conn

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

Let's analyze how ConnectionManager plays:

• He __enter__() The method is called when the execution enters the context of the with statement. It is responsible for configuring the context, in this case, connecting to a database. Returns the resource that needs to be managed: the connection to the database. Note that we have used the Optional write from typing module for connection object conn. We use Optional when the value can be one of two types: here a valid connection object or None.
• He __exit__() Method: Called when the execution leaves the context of the with statement. Handles the cleanup action of closing the connection. The parameters exc_type, exc_valueand traceback They are for handling exceptions within the `with` block. These can be used to determine if the context was exited due to an exception.

Now let's use the ConnectionManager in it with statement. We do the following:

• Try connecting to the database.
• Create a cursor to run queries
• Create a table and insert records.
• Query the database table and retrieve the query results.

db_name = "library.db"

# Using ConnectionManager context manager directly
with ConnectionManager(db_name) as conn:
	cursor = conn.cursor()

	# Create a books table if it doesn't exist
	cursor.execute("""
    	CREATE TABLE IF NOT EXISTS books (
        	id INTEGER PRIMARY KEY,
        	title TEXT,
        	author TEXT,
        	publication_year INTEGER
    	)
	""")

	# Insert sample book records
	books_data = (
    	("The Great Gatsby", "F. Scott Fitzgerald", 1925),
    	("To Kill a Mockingbird", "Harper Lee", 1960),
    	("1984", "George Orwell", 1949),
    	("Pride and Prejudice", "Jane Austen", 1813)
	)
	cursor.executemany("INSERT INTO books (title, author, publication_year) VALUES (?, ?, ?)", books_data)
	conn.commit()

	# Retrieve and print all book records
	cursor.execute("SELECT * FROM books")
	records = cursor.fetchall()
	print("Library Catalog:")
	for record in records:
    	    book_id, title, author, publication_year = record
    	    print(f"Book ID: {book_id}, Title: {title}, Author: {author}, Year: {publication_year}")
            cursor.close()

By running the above code you should get the following result:

Output >>>

Library Catalog:
Book ID: 1, Title: The Great Gatsby, Author: F. Scott Fitzgerald, Year: 1925
Book ID: 2, Title: To Kill a Mockingbird, Author: Harper Lee, Year: 1960
Book ID: 3, Title: 1984, Author: George Orwell, Year: 1949
Book ID: 4, Title: Pride and Prejudice, Author: Jane Austen, Year: 1813

2. Using contextlib's @contextmanager decorator

He contextlib The module provides the @contextmanager Decorator that can be used to define a builder function as a context manager. Here's how we do it for the database connection example:

# Writing a generator function with the `@contextmanager` decorator
import sqlite3
from contextlib import contextmanager

@contextmanager
def database_connection(db_name: str):
    conn = sqlite3.connect(db_name)
    try:
        yield conn  # Provide the connection to the 'with' block
    finally:
        conn.close()  # Close the connection upon exiting the 'with' block

This is how the database_connection the function works:

• He database_connection The function first establishes a connection that the yield The statement then provides the connection to the code block in the with declaration block. Please note that because yield In itself it is not immune to exceptions, we wrap it in a try block.
• He finally The block ensures that the connection is always closed whether an exception is thrown or not, ensuring that there are no resource leaks.

As we did before, let's use this in a way with statement:

db_name = "library.db"

# Using database_connection context manager directly
with database_connection(db_name) as conn:
	cursor = conn.cursor()

	# Insert a set of book records
	more_books_data = (
    	("The Catcher in the Rye", "J.D. Salinger", 1951),
    	("To the Lighthouse", "Virginia Woolf", 1927),
    	("Dune", "Frank Herbert", 1965),
    	("Slaughterhouse-Five", "Kurt Vonnegut", 1969)
	)
	cursor.executemany("INSERT INTO books (title, author, publication_year) VALUES (?, ?, ?)", more_books_data)
	conn.commit()

	# Retrieve and print all book records
	cursor.execute("SELECT * FROM books")
	records = cursor.fetchall()
	print("Updated Library Catalog:")
	for record in records:
    	    book_id, title, author, publication_year = record
    	    print(f"Book ID: {book_id}, Title: {title}, Author: {author}, Year: {publication_year}")
        cursor.close()

We connect to the database, insert some more records, query the database, and retrieve the query results. Here is the result:

Output >>>

Updated Library Catalog:
Book ID: 1, Title: The Great Gatsby, Author: F. Scott Fitzgerald, Year: 1925
Book ID: 2, Title: To Kill a Mockingbird, Author: Harper Lee, Year: 1960
Book ID: 3, Title: 1984, Author: George Orwell, Year: 1949
Book ID: 4, Title: Pride and Prejudice, Author: Jane Austen, Year: 1813
Book ID: 5, Title: The Catcher in the Rye, Author: J.D. Salinger, Year: 1951
Book ID: 6, Title: To the Lighthouse, Author: Virginia Woolf, Year: 1927
Book ID: 7, Title: Dune, Author: Frank Herbert, Year: 1965
Book ID: 8, Title: Slaughterhouse-Five, Author: Kurt Vonnegut, Year: 1969

Note that we open and close the cursor object. Then you can also use the cursor object in a with statement. I suggest trying it as a quick exercise!

Ending

And that's a wrap. I hope you learned how to create your own custom context managers. We look at two approaches: using a class with __enter__() and __exit()__ methods and using a generating function decorated with the @contextmanager decorator.

It's pretty easy to see that you get the following benefits from using a context manager:

• Resource setup and teardown are handled automatically, minimizing resource leaks.
• The code is cleaner and easier to maintain.
• Cleaner exception handling when working with resources.

As always, you can find the code on GitHub. Keep coding!

twitter.com/balawc27″ rel=”noopener”>Bala Priya C. is a developer and technical writer from India. He enjoys working at the intersection of mathematics, programming, data science, and content creation. His areas of interest and expertise include DevOps, data science, and natural language processing. He likes to read, write, code and drink coffee! Currently, he is working to learn and share his knowledge with the developer community by creating tutorials, how-to guides, opinion pieces, and more. Bala also creates engaging resource descriptions and coding tutorials.