Daily Python - Wtih Statements

By Maximus Meadowcroft | December 28, 2024

Context Managers and the `with` Statement

What Does the `with` Statement in Python Actually Do?

The with statement in Python is a clean and efficient way to manage resources. It ensures that setup and teardown operations are handled properly, such as opening and closing files, acquiring and releasing locks, or managing database connections. Python provides context managers to facilitate this resource management.

File Handling with `with`

The most common use of the with statement is for opening files in Python. It is the recommended approach because it ensures the file is properly closed after usage, even if an exception occurs.

with open("test.txt") as f:
    data = f.read()

Here, the with statement manages the file resource, eliminating the need for an explicit f.close() call. This is achieved through Python's context manager mechanism.

Creating a Context Manager

You can create a custom context manager by defining a class with two special methods:

__enter__: Sets up and returns the resource to be managed.
__exit__: Handles cleanup logic when exiting the with block.

Example

class CustomContextManager:  
    def __init__(self):  
        print("init called")  

    def __enter__(self):  
        print("enter __enter__")  
        return self  

    def __exit__(self, e_type, e_value, traceback):  
        print("exit __exit__")  

with CustomContextManager() as my_context:  
    print("with statement block")

Output:

init called
enter __enter__
with statement block
exit __exit__

Building a Custom File Manager

To create your own file manager, you need to manage the file's context explicitly and ensure proper cleanup. Here's an example:

class CustomFileManager:  
    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, e_type, e_value, traceback):  
        print("The file is closed :)")  
        self.file.close()  

with CustomFileManager('test.txt', 'w') as f:  
    f.write("Test")

Output

The file is closed :)

Use Cases for Context Managers

Context managers are often underutilized because many developers are unaware of their versatility. Below are several practical use cases:

1. Database Resource Management

Context managers can handle database connections efficiently, ensuring they are properly closed after use. Many libraries have built-in support for this.

with db_connection() as conn:
    conn.execute('SELECT * FROM table')

2. Thread Locks

For thread-safe operations, context managers can acquire and release locks, preventing deadlocks.

from threading import Lock

lock = Lock()
with lock:
    # Critical section
    shared_resource += 1

3. Temporary Environment Changes

Context managers can temporarily change the environment, such as switching directories or modifying environment variables.

import os

class ChangeDirectory:
    def __init__(self, new_path):
        self.new_path = new_path
        self.original_path = os.getcwd()

    def __enter__(self):
        os.chdir(self.new_path)

    def __exit__(self, e_type, e_value, traceback):
        os.chdir(self.original_path)

with ChangeDirectory('/tmp'):
    # Current working directory is now /tmp
    pass

4. Suppressing Specific Exceptions

The suppress context manager from the contextlib module can ignore specific exceptions.

from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove('non_existent_file.txt')

5. Timing Code Execution

A custom timer context manager can measure the execution time of a code block.

import time  

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

    def __exit__(self, e_type, e_value, traceback):  
        self.end = time.time()  
        Time = self.end - self.start  
        print(f"{Time = :.2f} seconds")  

with Timer():  
    time.sleep(2)