How To Create Customized Context Managers in Python

Context managers in Python allow you to work extra effectively with assets—facilitating setup and teardown of assets even when there are errors when working with the assets. Within the tutorial on writing environment friendly Python code, I coated what context managers are and why they’re useful. And in 3 Attention-grabbing Makes use of of Python’s Context Managers, I went over using context managers in managing subprocesses, database connections, and extra.

On this tutorial, you’ll discover ways to create your individual customized context managers. We’ll evaluate how context manages work after which have a look at the alternative ways you’ll be able to write your individual. Let’s get began.

What Are Context Managers in Python?

Context managers in Python are objects that allow the administration of assets resembling file operations, database connections, or community sockets inside a managed block of code. They be sure that assets are correctly initialized earlier than the block of code executes and routinely cleaned up afterward, no matter whether or not the code block completes usually or raises an exception.

Generally, context managers in Python have the next two particular strategies: __enter__() and __exit__(). These strategies outline the habits of the context supervisor when coming into and exiting a context.

How Do Context Managers Work?

When working with assets in Python, you must take into account organising the useful resource, anticipate errors, implement exception dealing with, and at last release the useful resource. To do that, you’ll most likely use a try-except-finally block like so:

strive: 
    # Establishing the useful resource
    # Working with the useful resource
besides ErrorType:
    # Deal with exceptions
lastly:
    # Unlock the useful resource

Primarily, we strive provisioning and dealing with the useful resource, apart from any errors which will come up throughout the course of, and at last release the useful resource. The lastly block is at all times executed no matter whether or not the operation succeeds or not. However with context managers and the with assertion, you’ll be able to have reusable try-except-finally blocks.

Now let’s go over how context managers work.

Enter Part (__enter__() technique):
When a with assertion is encountered, the __enter__() technique of the context supervisor is invoked. This technique is liable for initializing and organising the useful resource resembling opening a file, establishing a database connection, and the like. The worth returned by __enter__() (if any) is made obtainable to the context block after the `as` key phrase.

Execute the Block of Code:
As soon as the useful resource is ready up (after __enter__() is executed), the code block related to the with assertion is executed. That is the operation you need to carry out on the useful resource.

Exit Part (__exit__() technique):
After the code block completes execution—both usually or as a consequence of an exception—the __exit__() technique of the context supervisor known as. The __exit__() technique handles cleanup duties, resembling closing the assets. If an exception happens throughout the code block, details about the exception (sort, worth, traceback) is handed to __exit__() for error dealing with.

To sum up:

• Context managers present a technique to handle assets effectively by making certain that assets are correctly initialized and cleaned up.
• We use the with assertion to outline a context the place assets are managed.
• The __enter__() technique initializes the useful resource, and the __exit__() technique cleans up the useful resource after the context block completes.

Now that we all know how context managers work, let’s proceed to jot down a customized context supervisor for dealing with database connections.

Creating Customized Context Managers in Python

You’ll be able to write your individual context managers in Python utilizing one of many following two strategies:

1. Writing a category with __enter__() and __exit__() strategies.
2. Utilizing the contextlib module which offers the contextmanager decorator to jot down a context supervisor utilizing generator features.

1. Writing a Class with __enter__() and __exit__() Strategies

You’ll be able to outline a category that implements the 2 particular strategies: __enter__() and __exit__() that management useful resource setup and teardown respectively. Right here we write a ConnectionManager class that establishes a connection to an SQLite database and closes the database connection:

import sqlite3
from typing import Non-obligatory

# Writing a context supervisor class
class ConnectionManager:
    def __init__(self, db_name: str):
        self.db_name = db_name
        self.conn: Non-obligatory[sqlite3.Connection] = None

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

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

Let’s break down how the ConnectionManager works:

• The __enter__() technique known as when the execution enters the context of the with assertion. It’s liable for organising the context, on this case, connecting to a database. It returns the useful resource that must be managed: the database connection. Observe that we’ve used the Non-obligatory sort from the typing module for the connection object conn. We use Non-obligatory when the worth could be considered one of two sorts: right here a legitimate connection object or None.
• The __exit__() technique: It is known as when the execution leaves the context of the with assertion. It handles the cleanup motion of closing the connection. The parameters exc_type, exc_value, and traceback are for dealing with exceptions throughout the `with` block. These can be utilized to find out whether or not the context was exited as a consequence of an exception.

Now let’s use the ConnectionManager within the with assertion. We do the next:

• Attempt to connect with the database
• Create a cursor to run queries
• Create a desk and insert information
• Question the database desk and retrieve the outcomes of the question

db_name = "library.db"

# Utilizing ConnectionManager context supervisor immediately
with ConnectionManager(db_name) as conn:
	cursor = conn.cursor()

	# Create a books desk if it would not exist
	cursor.execute("""
    	CREATE TABLE IF NOT EXISTS books (
        	id INTEGER PRIMARY KEY,
        	title TEXT,
        	writer TEXT,
        	publication_year INTEGER
    	)
	""")

	# Insert pattern guide information
	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 guide information
	cursor.execute("SELECT * FROM books")
	information = cursor.fetchall()
	print("Library Catalog:")
	for document in information:
    	    book_id, title, writer, publication_year = document
    	    print(f"Book ID: {book_id}, Title: {title}, Author: {author}, Year: {publication_year}")
            cursor.shut()

Working the above code ought to provide the following output:

Output >>>

Library Catalog:
Ebook ID: 1, Title: The Nice Gatsby, Creator: F. Scott Fitzgerald, 12 months: 1925
Ebook ID: 2, Title: To Kill a Mockingbird, Creator: Harper Lee, 12 months: 1960
Ebook ID: 3, Title: 1984, Creator: George Orwell, 12 months: 1949
Ebook ID: 4, Title: Pleasure and Prejudice, Creator: Jane Austen, 12 months: 1813

2. Utilizing the @contextmanager Decorator From contextlib

The contextlib module offers the @contextmanager decorator which can be utilized to outline a generator perform as a context supervisor. Here is how we do it for the database connection instance:

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

@contextmanager
def database_connection(db_name: str):
    conn = sqlite3.join(db_name)
    strive:
        yield conn  # Present the connection to the 'with' block
    lastly:
        conn.shut()  # Shut the connection upon exiting the 'with' block

Right here’s how the database_connection perform works:

• The database_connection perform first establishes a connection which the yield assertion then provisions the connection to the block of code within the with assertion block. Observe that as a result of yield itself isn’t proof against exceptions, we wrap it in a strive block.
• The lastly block ensures that the connection is at all times closed, whether or not an exception was raised or not, making certain there are not any useful resource leaks.

Like we did beforehand, let’s use this in a with assertion:

db_name = "library.db"

# Utilizing database_connection context supervisor immediately
with database_connection(db_name) as conn:
	cursor = conn.cursor()

	# Insert a set of guide information
	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 guide information
	cursor.execute("SELECT * FROM books")
	information = cursor.fetchall()
	print("Updated Library Catalog:")
	for document in information:
    	    book_id, title, writer, publication_year = document
    	    print(f"Book ID: {book_id}, Title: {title}, Author: {author}, Year: {publication_year}")
        cursor.shut()

We hook up with the database, insert some extra information, question the db, and fetch the outcomes of the question. Right here’s the output:

Output >>>

Up to date Library Catalog:
Ebook ID: 1, Title: The Nice Gatsby, Creator: F. Scott Fitzgerald, 12 months: 1925
Ebook ID: 2, Title: To Kill a Mockingbird, Creator: Harper Lee, 12 months: 1960
Ebook ID: 3, Title: 1984, Creator: George Orwell, 12 months: 1949
Ebook ID: 4, Title: Pleasure and Prejudice, Creator: Jane Austen, 12 months: 1813
Ebook ID: 5, Title: The Catcher within the Rye, Creator: J.D. Salinger, 12 months: 1951
Ebook ID: 6, Title: To the Lighthouse, Creator: Virginia Woolf, 12 months: 1927
Ebook ID: 7, Title: Dune, Creator: Frank Herbert, 12 months: 1965
Ebook ID: 8, Title: Slaughterhouse-5, Creator: Kurt Vonnegut, 12 months: 1969

Observe that we open and shut the cursor object. So you can too use the cursor object in a with assertion. I counsel attempting that as a fast train!

Wrapping Up

And that’s a wrap. I hope you discovered how you can create your individual customized context managers. We checked out two approaches: utilizing a category with __enter__() and __exit()__ strategies and utilizing a generator perform adorned with the @contextmanager decorator.

It’s fairly simple to see that you just get the next benefits when utilizing a context supervisor:

• Setup and teardown of assets are routinely managed, minimizing useful resource leaks.
• The code is cleaner and simpler to take care of.
• Cleaner exception dealing with when working with assets.

As at all times, you’ll be able to discover the code on GitHub. Maintain coding!

Bala Priya C is a developer and technical author from India. She likes working on the intersection of math, programming, information science, and content material creation. Her areas of curiosity and experience embrace DevOps, information science, and pure language processing. She enjoys studying, writing, coding, and occasional! At the moment, she’s engaged on studying and sharing her data with the developer neighborhood by authoring tutorials, how-to guides, opinion items, and extra. Bala additionally creates partaking useful resource overviews and coding tutorials.