29 Skillful Use of Context Managers and With Statement for Concise Code

29 Skillful Use of Context Managers and With Statement for Concise Code #

Hello, I’m Jingxiao.

I believe you are familiar with the with statement in Python, as it has been mentioned multiple times in this column, especially in file input/output operations. However, most people may be accustomed to using it without knowing the “secret” behind it.

So, how exactly should we use the with statement? What is the context manager associated with it, and what is the relationship between them? In this lesson, I will unveil the mystery for you.

What is a Context Manager? #

In any programming language, common resource management operations include file input/output and database connection/disconnection. However, resources are limited, and it is essential to ensure that these resources are released after use to prevent resource leaks. Failure to do so can result in slow system processing or system crashes.

To better understand this concept, let’s look at the following example:

for x in range(10000000):
    f = open('test.txt', 'w')
    f.write('hello')

In this example, we opened 10,000,000 files without closing them afterward. If you run this code, you will encounter an error:

OSError: [Errno 23] Too many open files in system: 'test.txt'

This is a typical example of a resource leak. Opening too many files simultaneously consumes too many resources and causes the system to crash.

To address this problem, various programming languages have introduced different mechanisms. In Python, the corresponding solution is a Context Manager. A Context Manager helps automatically allocate and release resources. The most common way to use a Context Manager is with the with statement. Therefore, the correct way to write the code in the previous example is as follows:

for x in range(10000000):
    with open('test.txt', 'w') as f:
        f.write('hello')

In this updated code, each time the file "test.txt" is opened and 'hello' is written to it, the file is automatically closed, and the corresponding resources are released, preventing resource leaks. Additionally, the code using the with statement can also be expressed in the following form:

f = open('test.txt', 'w')
try:
    f.write('hello')
finally:
    f.close()

It is worth noting that the finally block is crucial, as it ensures that the file is closed, even if an error or exception occurs during the file write. However, compared to the with statement, this code appears to be more verbose and prone to omission, so we generally prefer using the with statement.

Another typical example is the threading.Lock class in Python. For example, if I want to acquire a lock, perform the corresponding operation, and then release the lock, the code can be written as follows:

some_lock = threading.Lock()
some_lock.acquire()
try:
    ...
finally:
    some_lock.release()

The equivalent code using the with statement is also concise:

some_lock = threading.Lock()
with some_lock:
    ...

From these two examples, we can see that the use of the with statement simplifies the code and effectively avoids resource leaks.

Implementation of Context Managers #

Class-based Context Managers #

After understanding the concept and benefits of context managers, let’s take a look at the implementation and internal workings of context managers through a specific example. Here, I have defined a custom context manager class called FileManager to simulate the file operations of opening and closing in Python:

class FileManager:
    def __init__(self, name, mode):
        print('calling __init__ method')
        self.name = name
        self.mode = mode 
        self.file = None
        
    def __enter__(self):
        print('calling __enter__ method')
        self.file = open(self.name, self.mode)
        return self.file


    def __exit__(self, exc_type, exc_val, exc_tb):
        print('calling __exit__ method')
        if self.file:
            self.file.close()
            
with FileManager('test.txt', 'w') as f:
    print('ready to write to file')
    f.write('hello world')

Note that when we create a context manager using a class, the class must include the __enter__() and __exit__() methods. In this case, the __enter__() method opens the file "test.txt" in write mode and returns the FileManager object assigned to the variable f. The __exit__() method is responsible for closing the previously opened file.

When we execute the context manager using the with statement:

with FileManager('test.txt', 'w') as f:
    f.write('hello world')

The following four steps occur in sequence:

  1. The __init__() method is called, initializing the FileManager object with the file name test.txt and file mode w.
  2. The __enter__() method is called, opening the file test.txt in write mode and returning the FileManager object.
  3. The string hello world is written to the file test.txt.
  4. The __exit__() method is called, responsible for closing the previously opened file.

Therefore, the output of this program is:

calling __init__ method
calling __enter__ method
ready to write to file
calling __exit__ method

Additionally, it is worth mentioning that the __exit__() method takes arguments exc_type, exc_val, exc_tb, which represent the exception type, exception value, and traceback, respectively. When we execute a with statement containing a context manager, if an exception is raised, the exception information will be included in these three variables and passed to the __exit__() method.

Thus, if you need to handle possible exceptions, you can add the corresponding code inside __exit__(), such as:

class Foo:
    def __init__(self):
        print('__init__ called')

    def __enter__(self):
        print('__enter__ called')
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        print('__exit__ called')
        if exc_type:
            print(f'exc_type: {exc_type}')
            print(f'exc_value: {exc_value}')
            print(f'exc_traceback: {exc_tb}')
            print('exception handled')
        return True

with Foo() as obj:
    raise Exception('exception raised').with_traceback(None)

# Output
__init__ called
__enter__ called
__exit__ called
exc_type: <class 'Exception'>
exc_value: exception raised
exc_traceback: <traceback object at 0x1046036c8>
exception handled

Here, we manually raise an exception "exception raised" within the with statement. You can see that the exception is successfully caught and handled in the __exit__() method. However, it is important to note that if __exit__() does not return True, the exception will still be raised. Therefore, if you are certain that the exception has been handled, add the statement return True at the end of __exit__().

Similarly, context managers are commonly used to represent database connection operations. Here is a simplified code for a database connection context manager:

class DBConnectionManager: 
    def __init__(self, hostname, port): 
        self.hostname = hostname 
        self.port = port 
        self.connection = None 
  
    def __enter__(self): 
        self.connection = DBClient(self.hostname, self.port) 
        return self 
  
    def __exit__(self, exc_type, exc_val, exc_tb): 
        self.connection.close() 
  
with DBConnectionManager('localhost', '8080') as db_client:

Similar to the previous example with FileManager:

  • The __init__() method is responsible for initializing the database, assigning the hostname and port to the variables hostname and port respectively.
  • The __enter__() method connects to the database and returns the DBConnectionManager object.
  • The __exit__() method is responsible for closing the database connection.

By doing so, every time you need to connect to the database in your program, you can simply use the with statement without worrying about closing the database or handling exceptions, greatly improving development efficiency.

Generator-based Context Managers #

Undoubtedly, class-based context managers are widely used and commonly seen in Python. However, Python’s context managers are not limited to class-based implementations. In addition to classes, context managers can also be implemented based on generators. Let’s take a look at an example.

For instance, you can use the decorator contextlib.contextmanager to define your own generator-based context manager that supports the with statement. Taking the previous example of the class-based context manager FileManager, we can represent it using the following form:

from contextlib import contextmanager

@contextmanager
def file_manager(name, mode):
    try:
        f = open(name, mode)
        yield f
    finally:
        f.close()

with file_manager('test.txt', 'w') as f:
    f.write('hello world')

In this code, the function file_manager() is a generator. When we execute the with statement, it opens the file and returns the file object f. After the with statement is executed, the closing file operation in the finally block will be performed.

You can see that when using generator-based context managers, there is no need to define __enter__() and __exit__() methods. However, it is essential to include the @contextmanager decorator, which is often overlooked by beginners.

After discussing the implementation of these two different types of context managers, it is important to emphasize that class-based and generator-based context managers are functionally equivalent. The difference lies in:

  • Class-based context managers are more flexible and suitable for large-scale system development.
  • Generator-based context managers are more convenient and concise, suitable for medium and small-scale programs.

Regardless of which type you use, do not forget to release resources in the __exit__() method or the finally block. This is particularly important.

Summary #

In this lesson, we first used a simple example to understand the ease of resource leakage and the serious consequences it can bring, and then introduced the concept of a solution - the context manager. The context manager is usually applied in scenarios such as opening and closing files and closing database connections, which ensures that used resources are quickly released and effectively improves the security of the program.

Next, we learned about the principles of context management through custom context manager examples, and together we learned about class-based context managers and generator-based context managers. Both of these have the same functionality, and the specific choice depends on your specific use case.

In addition, context managers are usually used together with the with statement, greatly improving the conciseness of the program. It is important to note that when we use the with statement to execute operations of the context manager, if an exception is thrown, the specific information such as the type and value of the exception will be passed as arguments to the __exit__() function. You can define relevant operations to handle the exception, and after handling the exception, don’t forget to add the statement return True, otherwise the exception will still be thrown.

Reflection Question #

So, in your daily learning and work, in which scenarios have you used context managers? What problems have you encountered or what new discoveries have you made during the process? Feel free to leave a comment below and discuss with me. Also, feel free to share this article so that we can exchange ideas and make progress together.