Context Managers - CameronAuler/python-devops GitHub Wiki

Context managers simplify resource management by ensuring that resources like files, database connections, and network sockets are properly opened and closed. They are typically used with the with statement, which ensures that cleanup (like closing a file) always happens, even if an error occurs.

Table of Contents

with Statement (Built-in Context Managers)

The with statement ensures automatic cleanup of resources (files, database connections, etc.). It avoids manual close() calls, reducing the chance of resource leaks and handles exceptions gracefully.

Example

If an exception occurs before file.close() is called, the file stays open. If we forget to close the file manually, it can lead to memory leaks.

file = open("example.txt", "r")  # Open a file
content = file.read()
print(content)
file.close()  # We must explicitly close the file

with (Recommended)

The with statement automatically closes the file, even if an error occurs. It provides automatic cleanup (no need for file.close()) and it is exception-safe (cleanup occurs even if an error happens inside the block).

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

Custom Context Managers

Class with __enter__() and __exit__()

To create a custom context manager:

  • Implement __enter__() → Runs when entering the with block. __enter__() opens the file and returns it.
  • Implement __exit__() → Runs when exiting the with block, ensuring cleanup. __exit__() closes the file when leaving the with block, even if an error occurs.
class FileManager:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file  # Returns the resource to be used inside `with`

    def __exit__(self, exc_type, exc_value, traceback):
        self.file.close()  # Ensures file is closed
        print("File closed.")  # Confirms cleanup

# Using the custom context manager
with FileManager("example.txt", "w") as f:
    f.write("Hello, World!")

# After `with`, the file is closed automatically
# Output:
File closed.

contextlib

contextlib.contextmanager (Simplified Context Manager)

Instead of creating a class, we can use a function-based approach with contextlib.contextmanager. The advantages of this include less boilerplate code (compared to defining __enter__ and __exit__ methods) and it works well for simple one-resource management cases.

from contextlib import contextmanager

@contextmanager
def open_file(filename, mode):
    file = open(filename, mode)
    try:
        yield file  # Provide file to `with` block
    finally:
        file.close()  # Ensures file is closed
        print("File closed.")

# Using the context manager
with open_file("example.txt", "w") as f:
    f.write("Hello, Python!")
# Output:
File closed.

Handling Context Manager Exceptions

If an exception occurs inside the with block, __exit__() will handle it. If an exception occurs, __exit__() captures it and prevents it from crashing the program. Returning True suppresses the exception. If False is returned (or nothing is returned), the error propagates.

class ErrorHandlingContext:
    def __enter__(self):
        print("Entering the context")
        return self  # Returns an object that can be used inside `with`

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            print(f"Exception caught: {exc_value}")
            return True  # Suppresses the exception
        print("Exiting normally.")

# Example with an exception
with ErrorHandlingContext():
    print("Inside the block")
    raise ValueError("An error occurred!")  # Exception is caught
# Output:
Entering the context
Inside the block
Exception caught: An error occurred!