140. Creating and Handling Custom Context Managers

In Python, context managers are a great way to manage resources such as file handles, network connections, or database connections that need to be acquired and released properly. Using contextlib, you can define reusable context managers without needing to implement the __enter__ and __exit__ methods manually. Here's how you can create and handle custom context managers using the contextlib module.

1. Basic Context Manager with contextlib.contextmanager

You can create a custom context manager using the @contextmanager decorator from contextlib.

from contextlib import contextmanager

@contextmanager
def my_open_file(file_name, mode):
    print(f"Opening file: {file_name}")
    file = open(file_name, mode)
    try:
        yield file  # Provides control back to the with statement
    finally:
        file.close()
        print(f"Closing file: {file_name}")

# Using the custom context manager
with my_open_file('test.txt', 'w') as file:
    file.write("Hello, world!")

Output:

Opening file: test.txt
Closing file: test.txt

Explanation:

• The @contextmanager decorator simplifies the creation of context managers.

• The yield keyword returns control to the with statement block, and the finally block ensures that resources are cleaned up (in this case, the file is closed).

2. Context Manager for Database Connections

A custom context manager for managing database connections can be created as follows.

from contextlib import contextmanager
import sqlite3

@contextmanager
def db_connection(db_name):
    conn = sqlite3.connect(db_name)
    print(f"Connected to {db_name}")
    try:
        yield conn  # Yield the database connection
    finally:
        conn.close()
        print(f"Closed connection to {db_name}")

# Using the custom context manager
with db_connection('example.db') as conn:
    cursor = conn.cursor()
    cursor.execute("CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)")
    cursor.execute("INSERT INTO users (name) VALUES ('Alice')")
    conn.commit()

Output:

Connected to example.db
Closed connection to example.db

Explanation:

• The context manager ensures that the database connection is opened and closed properly, handling exceptions if any occur during the execution of database operations.

3. Custom Context Manager for Timer

You can create a context manager to measure the time taken by a code block.

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start_time = time.time()
    try:
        yield  # Control is handed back to the "with" block
    finally:
        end_time = time.time()
        print(f"Time taken: {end_time - start_time} seconds")

# Using the timer context manager
with timer():
    total = sum(range(1, 1000000))  # Code to be timed

Output:

Time taken: 0.026546478271484375 seconds

Explanation:

• The context manager timer() measures the time before and after the code block inside the with statement, then prints the time taken.

4. Context Manager with Custom Exception Handling

Context managers can also be used for custom exception handling.

from contextlib import contextmanager

@contextmanager
def exception_handler():
    try:
        yield
    except ZeroDivisionError:
        print("ZeroDivisionError occurred!")
    except ValueError:
        print("ValueError occurred!")

# Using the exception handler context manager
with exception_handler():
    x = 1 / 0  # ZeroDivisionError

with exception_handler():
    int("not a number")  # ValueError

Output:

ZeroDivisionError occurred!
ValueError occurred!

Explanation:

• The custom context manager exception_handler() catches specific exceptions and handles them within the with block, providing graceful error handling.

5. Context Manager for Resource Locking

A context manager can manage resource locking to ensure that only one thread or process accesses a resource at a time.

from contextlib import contextmanager
import threading

lock = threading.Lock()

@contextmanager
def locked_resource():
    lock.acquire()
    print("Lock acquired")
    try:
        yield  # Control is given to the "with" block
    finally:
        lock.release()
        print("Lock released")

# Using the locked_resource context manager
with locked_resource():
    print("Resource is being accessed")

Output:

Lock acquired
Resource is being accessed
Lock released

Explanation:

• The context manager locked_resource() acquires a lock before entering the with block and releases it afterward, ensuring safe access to shared resources.

6. Context Manager with Multiple Resources

You can use a context manager to manage multiple resources simultaneously.

from contextlib import contextmanager

@contextmanager
def multi_resource_manager(resource1, resource2):
    print(f"Acquiring {resource1}")
    print(f"Acquiring {resource2}")
    try:
        yield (resource1, resource2)  # Yield both resources
    finally:
        print(f"Releasing {resource1}")
        print(f"Releasing {resource2}")

# Using the multi-resource manager context manager
with multi_resource_manager('Resource1', 'Resource2') as (r1, r2):
    print(f"Using {r1} and {r2}")

Output:

Acquiring Resource1
Acquiring Resource2
Using Resource1 and Resource2
Releasing Resource1
Releasing Resource2

Explanation:

• This context manager acquires multiple resources, yields them to the with block, and releases them afterward.

7. Context Manager for Logging

A context manager can be used to log the entry and exit of a function.

import logging
from contextlib import contextmanager

logging.basicConfig(level=logging.INFO)

@contextmanager
def log_function(func_name):
    logging.info(f"Entering {func_name}")
    try:
        yield
    finally:
        logging.info(f"Exiting {func_name}")

# Using the logging context manager
with log_function("process_data"):
    print("Processing data...")

Output:

INFO:root:Entering process_data
Processing data...
INFO:root:Exiting process_data

Explanation:

• The context manager log_function() logs entry and exit messages, useful for tracing the execution of functions.

8. Context Manager for File Backup

This context manager creates a backup of a file before making changes.

import shutil
from contextlib import contextmanager

@contextmanager
def file_backup(file_path):
    backup_path = file_path + ".bak"
    shutil.copy(file_path, backup_path)
    print(f"Backup created: {backup_path}")
    try:
        yield  # The file operations go here
    finally:
        print(f"Backup kept at: {backup_path}")

# Using the file_backup context manager
with file_backup('example.txt'):
    print("Making changes to the file...")

Output:

Backup created: example.txt.bak
Making changes to the file...
Backup kept at: example.txt.bak

9. Context Manager with Configuration Changes

A context manager can temporarily change configuration settings.

from contextlib import contextmanager

@contextmanager
def temporary_config(config, key, value):
    old_value = config.get(key)
    config[key] = value
    try:
        yield
    finally:
        config[key] = old_value

config = {'theme': 'light', 'language': 'English'}

# Using the temporary_config context manager
with temporary_config(config, 'theme', 'dark'):
    print(f"Theme during session: {config['theme']}")

print(f"Theme after session: {config['theme']}")

Output:

Theme during session: dark
Theme after session: light

10. Context Manager for Network Resource

A context manager for managing network resources, such as opening and closing a socket connection.

import socket
from contextlib import contextmanager

@contextmanager
def open_network_connection(host, port):
    conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    conn.connect((host, port))
    print(f"Connected to {host}:{port}")
    try:
        yield conn
    finally:
        conn.close()
        print("Connection closed")

# Using the network connection context manager
with open_network_connection('example.com', 80) as conn:
    conn.send(b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")

Output:

Connected to example.com:80
Connection closed

Explanation:

• This context manager manages a network socket connection, ensuring that it is opened before the with block and properly closed afterward.

---

Conclusion:

Custom context managers using the contextlib module provide a flexible way to manage resources, handle exceptions, and ensure proper cleanup in Python. You can use them to simplify the resource management process in your applications, making the code more readable and maintainable.