Python Context Manager

Introduction

While Python's context managers are widely used, few understand the purpose behind their use. These statements, commonly used with reading and writing files, assist the application in conserving system memory and improve resource management by ensuring specific resources are only in use for certain processes.

In this lab, you will learn how to use the with statement for automatic resource management, and then you will create your own context managers using two different methods: a function-based approach and a class-based approach.

The with statement

In this step, you will learn how to use the with statement with context managers like file objects. A context manager is an object that is notified when a context (a block of code) starts and ends. The with statement takes care of this notification process automatically.

For example, file objects are context managers. When a with block is finished, the file object is closed automatically, even if errors occur within the block. This prevents data loss and resource leaks.

Note: If you are not familiar with using WebIDE (VS Code like), please learn it first at this lab.

Let's see this in action. In the file explorer on the left, open the main.py file. Copy and paste the following code into it. This code will first create and write to a file, then read its contents, all using the with statement.

## We will create a file to demonstrate the 'with' statement.
## First, let's write some content to a file.
with open('my_file.txt', 'w') as f:
    f.write('This is a test file.')

## Now, let's read it back using a context manager.
## The file is automatically closed when exiting the 'with' block.
with open('my_file.txt') as f:
    file_contents = f.read()
    print(file_contents)

## The file is automatically closed here, even if an error occurred.
print("\nFile has been read and closed.")

Now, run the script from the terminal to see the output.

python3 main.py

You should see the following output, confirming that the file was written to and read from successfully. The with statement ensures that my_file.txt is properly closed after both the write and read operations.

This is a test file.

File has been read and closed.

Using a with statement is the recommended way to handle resources like files or network connections because it makes your code cleaner and safer.

Writing your own context manager

In this step, you will discover how to create a custom context manager using a generator function and the @contextlib.contextmanager decorator. This is a concise way to create simple context managers.

The code before the yield statement runs when entering the with block (the setup logic), and the code after the yield statement runs when exiting it (the cleanup logic). The value that is yielded can be assigned to a variable using the as keyword in the with statement.

Open the function_cm.py file from the file explorer. Copy the following code into the file.

## Function-based context manager using contextlib decorator
import contextlib

@contextlib.contextmanager
def context_manager(num):
    print('Enter')  ## Code before yield runs on __enter__
    yield num + 1   ## Value yielded becomes 'cm' variable
    print('Exit')    ## Code after yield runs on __exit__

with context_manager(2) as cm:  ## cm receives the yielded value (3)
    print('Right in the middle with cm = {}'.format(cm))

Now, execute the script in the terminal.

python3 function_cm.py

You will see the following output. Notice the order of execution: the code before yield runs first, then the code inside the with block, and finally the code after yield.

Enter
Right in the middle with cm = 3
Exit

This approach is very useful for creating temporary contexts, such as setting up and tearing down a database connection or temporarily changing a setting.

Class based context manager

In this step, you will implement a context manager as a class. This approach provides more control and is suitable for more complex scenarios. To create a class-based context manager, you must define two special methods: __enter__ and __exit__.

• __enter__: This method is called when the with statement is entered. It should perform the setup logic. Its return value is assigned to the variable after as, if provided.
• __exit__: This method is called when the with block is exited, either normally or due to an exception. It should perform the cleanup logic.

Let's create one. Open the class_cm.py file and copy the following code into it.

## Class-based context manager: implement __enter__ and __exit__ methods
class ContextManager:
    def __enter__(self, *args, **kwargs):  ## Called when entering 'with' block
        print("--enter--")
        return self  ## Can return object to use as 'as' variable

    def __exit__(self, *args):  ## Called when exiting 'with' block
        print("--exit--")

with ContextManager():  ## Calls __enter__, then __exit__ when done
    print("test")

Run the script from the terminal.

python3 class_cm.py

The output demonstrates the flow of control. The __enter__ method is called first, followed by the code inside the with block, and finally, the __exit__ method is called to clean up.

--enter--
test
--exit--

The class-based approach is powerful because it allows you to encapsulate complex setup and teardown logic within a single, reusable object.

Summary

Congratulations on completing the lab! You have learned how to effectively manage resources in Python using context managers.

You started with the common with statement for safe file handling. Then, you moved on to creating your own custom context managers using two different methods:

• A function-based approach with the @contextlib.contextmanager decorator, which is great for simple setup and teardown tasks.
• A class-based approach by implementing the __enter__ and __exit__ methods, which is ideal for more complex resource management.

These skills are crucial for writing clean, reliable, and efficient Python code that properly handles resources.