Creating Custom Context Managers in Python (66/100 Days of Python)

Day 66 of the “100 Days of Python” blog post series covering how to create custom context managers

Python’s context managers provide a convenient way to manage resources in a Python program, such as file handles or database connections. Context managers help ensure that resources are properly managed and released after use, even if an error occurs.

What are Context Managers?

Before we dive into creating custom context managers, let’s review what context managers are and how they work in Python.

Context managers are objects that define a setup and a teardown method. The setup method is called when the context is entered, and the teardown method is called when the context is exited. This ensures that resources are properly managed and released, even if an error occurs.

In Python, the with statement is used to create a context. The with statement takes an object that acts as a context manager and a code block to be executed. For example, here's how we might use the built-in open() function as a context manager to ensure that a file is properly closed after use:

with open('file.txt', 'r') as f:
    data = f.read()
    # do something with data

In this example, open('file.txt', 'r') returns a file object, which is a context manager. When the with block is entered, the __enter__() method of the file object is called, and the file is opened. When the block is exited, the __exit__() method is called, and the file is closed.

Creating Custom Class-Based Context Managers

The most common way to create a custom context manager in Python is by using a class. To create a class-based context manager, we need to define two methods: __enter__() and __exit__():

class MyContextManager:
    def __enter__(self):
        # Code to setup context
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        # Code to teardown context
        pass

In this example, the __enter__() method is called when the context is entered, and the __exit__() method is called when the context is exited. The __enter__() method should return an object that can be used inside the with block. In this case, we simply return self. Let’s look at an example of how we might use this context manager:

with MyContextManager() as my_context:
    # Code to use the context

In this example, MyContextManager() creates a new instance of the context manager class. When the with block is entered, the __enter__() method of the context manager is called, and the returned object (self) is assigned to the variable my_context. The code inside the with block can then use my_context as needed. When the block is exited, the __exit__() method is called.

Creating Custom Function-Based Context Managers

In addition to class-based context managers, Python also allows us to create function-based context managers using the contextlib module. Function-based context managers are simpler and more concise than class-based context managers, but they have some limitations.

To create a function-based context manager, we can use the contextlib.contextmanager decorator:

from contextlib import contextmanager


@contextmanager
def my_context_manager():
    # Code to setup context
    try:
        yield
    finally:
        # Code to teardown context

In this example, the @contextmanager decorator is used to create a function-based context manager. The function should include a yield statement, which defines the point at which the code inside the with block will execute. The code before the yield statement is used to set up the context, and the code after the yield statement is used to tear down the context. The yield statement itself acts as the point where the code inside the with block is executed:

with my_context_manager():
    # Code to use the context

When the with block is entered, the code before the yield statement is executed to set up the context. The yield statement acts as a placeholder for the code inside the with block, which is executed at this point. When the block is exited, the code after the yield statement is executed to tear down the context.

Advantages of Custom Context Managers

Custom context managers can be useful in a variety of situations:

• Resource management: Custom context managers can be used to manage resources such as file handles, database connections, and network sockets. By using a context manager, we can ensure that these resources are properly managed and released, even if an error occurs.
• Locking: Custom context managers can be used to implement locking mechanisms. For example, we might create a context manager that acquires a lock before the code inside the with block is executed, and releases the lock when the block is exited.
• Performance monitoring: Custom context managers can be used to implement performance monitoring. For example, we might create a context manager that measures the execution time of the code inside the with block and logs it to a file or database.
• Error handling: Custom context managers can be used to implement error handling. For example, we might create a context manager that catches exceptions and logs them to a file or database.

Context managers provide a powerful and flexible way to manage resources and implement common programming patterns in Python. By creating our own custom context managers, we can tailor these patterns to our specific needs and simplify our code.

What’s next?

• If you found this story valuable, please consider clapping multiple times (this really helps a lot!)
• Hands-on Practice: Free Python Course
• Full series: 100 Days of Python
• Previous topic: Pathlib
• Next topic: All the Ways You Can Use Context Managers