Understanding Context Manager in Python

Python is a versatile and powerful programming language, known for its simplicity and readability. One of the features that contribute to its elegance is the context manager. Context managers are a way to manage resources efficiently, ensuring that they are properly acquired and released, even in the face of exceptions. In this article, we’ll explore what context managers are, how they work, and how you can create your own custom context managers.

Introduction to Context Managers

A context manager in Python is an object that defines the runtime context to be established when executing a block of code. It is primarily used for resource management, ensuring that resources are properly acquired and released. Common use cases include:

Opening and closing files.
Acquiring and releasing locks.
Establishing and closing database connections.
Managing temporary changes to the environment.

Context managers are implemented using the with statement, which simplifies resource management and makes the code cleaner and more readable.

It is very similar to the "Using" keyword in C#

The with Statement

The with statement is used to wrap the execution of a block of code within methods defined by a context manager. The syntax is as follows:

with context_manager as resource:
    # Code block

When the with block is entered, the context manager’s __enter__ method is called, and when the block is exited, the __exit__ method is called. This ensures that resources are properly managed, even if an exception occurs within the block.

Method	Description
`.__enter__(self)`	This method handles the setup logic and is called when entering a new `with` context. Its return value is bound to the `with` target variable.
`.__exit__(self, exc_type, exc_value, exc_tb)`	This method handles the teardown logic and is called when the flow of execution leaves the `with` context. If an exception occurs, then `exc_type`, `exc_value`, and `exc_tb` hold the exception type, value, and traceback information, respectively.

Built-in Context Managers

Python provides several built-in context managers. Some of the most commonly used ones include:

File Handling

with open("example.txt", "r") as file:
    content = file.read()
    # File is automatically closed when the block is exited

Creating Custom Context Managers

You can create your own context managers by defining a class with __enter__ and __exit__ methods or by using the contextlib module.

Using Classes

Here’s an example of a custom context manager that creates a backup of a file before allowing modifications:

# Example: File Backup Context Manager
import os
import shutil

class FileBackupContextManager:
    def __init__(self, file_path):
        self.file_path = file_path
        self.backup_file_path = f'{file_path}_backup.txt'
    

    def __enter__(self):
        # if the file path exists, the make a copy in backup file
        if os.path.exists(self.file_path):
            shutil.copy2(self.file_path, self.backup_file_path)
        return self
    

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            print(f"An error occured: {exc_value}. Restore the original file")
            if os.path.exists(self.backup_file_path):
                shutil.copy2(self.backup_file_path, self.file_path)
        else:
            # no exception, delete backup file
            if os.path.exists(self.backup_file_path):
                os.remove(self.backup_file_path)


file_path = 'sample.txt'

with open(file_path, 'w') as file:
    file.write("Some more context")

def class_context_manager_example():

    try:
        with FileBackupContextManager(file_path=file_path) as backup_manager:
            with open(file_path, 'w') as file:
                file.write("This is new content")

                # mock error
                raise ValueError("Something went wrong")
    except Exception as ex:
        print(f'Error: {ex}')


class_context_manager_example()

# read file content
with open(file_path, 'r') as file:
    print(f"File content after changes: {file.read()}")

The above context manager creates a backup of file before making any changes to it. In case of any error it restored the file that was backed up

If the file modification is successfull, it makes sure to delete the backup file

Using Function

Below is the same custom context manager using a function. It server the same purpose but the structure is different.

In the function context manager we use "yield" keyword to make the scope code run. So, wherever you see "yield" consider it as the code written inside the "with" scope

# Example: File Backup Context Manager

from contextlib import contextmanager
import os
import shutil


def take_file_back(file_path, backup_file_path):
    """
    Take file backup by coping main file to backupfile
    """
    if os.path.exists(file_path):
        shutil.copy2(file_path, backup_file_path)


def restore_backup_file(file_path, backup_file_path):
    """
    Copy backup file to main file
    """
    print(f"An error occured. Restore the original file")
    if os.path.exists(backup_file_path):
        shutil.copy2(backup_file_path, file_path)


def delete_backup_file(backup_file_path):
    """
    Delete the backup file, as it is no longer required
    """
    if os.path.exists(backup_file_path):
        os.remove(backup_file_path)


@contextmanager
def file_backup_context_manager(file_path):
    backup_file_path = f'{file_path}_backup.txt'
    
    # take file_backup
    take_file_back(file_path, backup_file_path)
    try:
        yield
        delete_backup_file(backup_file_path)
    except Exception as ex:
        restore_backup_file(file_path, backup_file_path)




file_path = 'sample.txt'

with open(file_path, 'w') as file:
    file.write("Some more context")


def function_context_manager_example():
    try:
        with file_backup_context_manager(file_path=file_path):
            with open(file_path, 'w') as file:
                file.write("This is new context for testing function context manager")

                # mock error
                raise ValueError("Something went wrong")
    
    except Exception as ex:
        print(f"Error: {ex}")


function_context_manager_example()

# read file content
with open(file_path, 'r') as file:
    print(f"File content after changes: {file.read()}")

Best Practices

Ensure Resource Cleanup: Always release resources in the __exit__ method.
Handle Exceptions Gracefully: Use the __exit__ method to handle exceptions and ensure resources are properly released.
Use contextlib for Simplicity: For simple context managers, consider using contextlib.contextmanager to reduce boilerplate code.

The code is part of my ongoing git repository where I push small logic related to Python.

Github Link for reference: https://github.com/vipulm124/python-concepts.git