Skip to content

[Python] How To Use @contextmanager Decorator

Last Updated on 2024-09-12 by Clay

In Python, the context manager decorator @contextmanager from the contextlib module allows developers to conveniently create our own context manager.

What is a context manager? Simply put, it defines a block that handles certain logic, and it automatically executes some code when 'entering' and 'exiting' that block. In Python, the most common usage can be seen with the with syntax, which automatically handles operations like opening/closing files or allocating/releasing resources.

Context Manager Without @contextmanager

If we want to create a context manager for managing files, the traditional approach would be to define a class and then define __enter__() and __exit__() methods to control the logic when entering and exiting.

from typing import Any


class FileManager:
    def __init__(self, filename: str, mode: str) -> None:
        self.file = None
        self.filename = filename
        self.mode = mode

    def __enter__(self) -> Any:
        print("Open file.")
        self.file = open(file=self.filename, mode=self.mode)
        return self.file

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        print("Close file.")
        if self.file:
            self.file.close()


if __name__ == "__main__":
    with FileManager("test.log", "w") as f:
        f.write("This is a test log.")


Output:

Open file.
Close file.


Of course, we can check if the file was successfully written. In short, our defined FileManager class will open the file and return the file object self.file when entering the with block, and automatically close the file when exiting the block.

This is an effective and traditional way to do it, but let's now look at how @contextmanager simplifies this process.

Using @contextmanager

@contextmanager is meant to simplify the method we just defined.

from contextlib import contextmanager


@contextmanager
def open_file(filename: str, mode: str):
    print("Open file.")

    f = open(file=filename, mode=mode)

    try:    
        yield f
    finally:
        print("Close file.")
        f.close


if __name__ == "__main__":
    with open_file("test.log", "r") as f:
        print(f.read())


Output:

Open file.
This is a test log.
Close file.


The code does exactly the same thing as the previous class-based approach, but with a simplified syntax. In short, @contextmanager essentially creates a generator. This generator pauses when it reaches the yield statement, returning the yield value to the variable following the with statement (in our example, this is f). Upon exiting, the context manager will always execute the code after the yield, which makes it perfect for try...finally blocks.

Here are a few scenarios where context managers are particularly useful:

  1. Releasing resources: In processes where resources must be released after completion, context managers are perfect for automatic exit handling.
  2. Temporarily changing variable values: Suppose you have a scenario where a variable needs to be temporarily changed during an operation, and then restored afterward. This can be handled using a context manager. (A recent use case I encountered was when updating a database — during the update, my system needed to be in maintenance mode, but even if there were errors in the user's data, the system had to return to online status afterward.)
  3. Accessing and releasing locks


To summarize, the usage of @contextmanager can be understood as follows:

  • @contextmanager is used to create simplified context managers, defining what to do when entering and exiting the context using yield.
  • When entering the context, the code before the yield is executed; upon exiting, the code after the yield is executed, usually to clean up.
  • Using it with try and finally ensures that resources are released properly, even in case of exceptions.

References

Read More

[Python] @property 裝飾器將類別方法轉為只讀屬性

Tags:

Leave a Reply