Context Managers in Python — Go Beyond “with open() as file”
Have you ever used the with statement elsewhere?
Introduction
When we deal with files in Python, the most common operation is probably the use of the built-in open() function. This creates a file object, which allows us to read and write data as applicable. When we use the open() function, we almost always use the function together with the with statement, because we learned its usage either from the official reference for the open function or some online tutorials. The basic form is shown below.
By running the above code, you’ll see that a file called hello.txt has been created in your current work directory. To verify that the Hello World! string has been written to the file, we can open the file to read the data:
As you can see above, we’re able to read the file, opening it with the open() function by specifying the read mode (r, as opposed to the w that we used previously for writing purposes).
Closing File Automatically
Many of us probably know why we use the with statement when we open a file here. For those who don’t know about it, please see the following code first:
In the code above, we modified the file by appending extra string to the file (note that we’re using the a mode for appending purposes). When we were done with the with statement, we found out that the file was closed, even though we didn’t explicitly call the close() method on the file object. This is exactly what the with statement did for us — it closed the file automatically when we exited the with statement.
But you may be wondering why it’s such a big deal that it closes the file for us. Consider the following trivial example:
As shown, we first modified the file by appending some new data. But then we forgot to close the file after this operation. When we read the file again, we couldn’t see the change that we thought we’d made, which could cause unexpected bugs in our code. As shown previously, if we had used the with statement, every file operation will be cleaned up by Python, closing the file for us automatically. More importantly, imagine that we can have more complicated operations with the file and some operations may involve exceptions that make the program stop running. In these possible scenarios, the file will still get a chance to be closed safely and automatically because of the use of the with statement.
Context Managers
In a broader sense, the with statement for file opening is an example of a context manager in use.
What’s a context manager? It’s a Python object that does the housekeeping for you when you use particular resources. Specifically, the context manager sets up a temporary context for you and destructs the context after all the operations are completed.
In terms of the file opening operation, what the context manager does can be demonstrated with try, except, and finally statements. Consider the following pseudocode for the probable implementation of the with statement:
As shown above, the context manager opens the file for you and creates a file object, which will be further manipulated. When we’ve finished the operation and any exceptions that are raised during the operation, the context manager will close the file for us. As shown above, because the files are shared resources and it’s your responsibility (certainly with the aid of context managers to make things easier), it’s critical that you releasing them when you’re done with your operation so other processes can access them.
Usage With Threading Management
As discussed, the with statement is best used when we need to work with something that is shared. One such usage is to deal with data in a multi-threaded projects. As you may know, when multi threads have access to the same pool of data, things can get messy. Something like what I showed with the data appending with the file operations can happen more easily. One operation is appending data, while the other is reading the old data in another file operation.
For example, one thread is trying to add data to a dictionary, while the other thread is trying to iterate the dictionary. Everything will easily go out of your control very soon. To address this problem, we can use a thread lock to help us mitigate the data mess. Importantly, because we want to have a full control of the resources in a temporary manner, it’s the best use case of the with statement. Let’s consider the following code for a trivial example.
As you can see, using the with statement can dramatically improve the conciseness of our code. More importantly, it will automatically release the lock when the operation is completed in the with statement. Without using the context manager (i.e., the with statement), we’ll have to manage these resources manually and carefully. If we forget to release a lock, our program will run into unexpected problems.
The Context Management Protocol
To manage some resources ourselves, we can create a context manager ourselves. One way to do this is to implement the methods for the context management protocol. You can conceptualize it as a duck typing — we’ll simply define the __enter__ and __exit__ magic methods without formally conforming the protocol or implementing the interface as you may do that in other programming languages. The following code shows you a proof of concept of how we do that:
As shown above, we simply defined a class, which implemented the __enter__ and __exit__ methods, which will make the instance of this class able to manage context for us. From the syntactical perspective, we can use this class in the with statement, as shown in Line 12. The printed texts clears show us the order of how these operations are coordinated nicely. Specifically, the created instance (as printed in line 15) would call the __enter__ method (as printed in line 16) to start the context, then we ran the operations ourselves (as printed in line 17), and finally, the context manager would exit the management by calling the __exit__ method.
The contextlib Module
You may find that it’s kind of tedious to implement the special methods __enter__ and __exit__ to create a context manager ourselves. With the contextlib module in the standard Python library, it’s much easier for context management. A full review of the entire module is beyond the scope of this article, I’ll just focus on a particular method in the module to create a context manager. But before we do that, let’s take a step back, and review decorators first, because the decorators technique is relevant here.
Decorators are functions that modify the behaviors of other functions without affecting their central functionalities. In other words, the decorated function will do whatever it’s supposed to do, but the decorator adds some flavors to it. We can think of the following example for a quick refresh of the decorator concept:
For decorators, you’ll just need to create a function that accepts another function as its input. The decoration is the operations defined in the decorator function. In this case, we’ll just simply make records before and after the function call. To use the decorator, we just placed the decorator function name with an @ sign prefix. As you can say, calling the decorated function (line 15) successfully resulted in extra logging before and after the function invocation.
With a basic understanding of decorators, let’s look at an example using the contextlib module to help us manage contexts in the following code snippet:
We use the decorator function contextmanager to decorate the context_manager_example function. In the body of the function, you may notice something unusual here — the yield keyword. You’ve probably seen the yield keyword when you’re learning generators, which are iterators that render the elements when they’re asked to do so (called lazy evaluation). In these use cases, yield means produce. You can learn more about generators in my previous articles (here and there).
Besides the meaning of produce, yield can also mean give way to, which is exactly happening during the context management decoration. Specifically, once the context manager (the decorated function context_manager_example) completes the setup, it yields execution, which allows the code in the with statement to run. After the operation is completed, it regains the control. Importantly, yield in Python is specially handled in Python such that we it has the control back, it will run from where it was yielded. This is why the print function following the yield keyword got called only once and right after the completion of the operations in the with statement.
Takeaways
In this article, we reviewed the concept of context managers through the example of the file operation involving the use of the with statement. We understood that it was the context manager that helped us do the housekeeping job by closing the file.
In a broader sense, context managers are useful to manage resources that are intended to be shared within your program or other programs in the computer. Context managers help us manage the acquisition and release of these shared resources responsibly. We also reviewed how we can override the __enter__ and __exit__ methods to create a custom context manager class. Alternatively, we can take advantage of the contextlib module to create context managers using decorators.
However, as you may notice, there are things that we did not cover. For example, the function signature of the __exit__ method has other parameters that we didn’t implement, such as exception handling. A full implementation of these parameters should be a case-by-case evaluation. If you’re interested, you can learn more advanced knowledge by referring to some more realistic examples, such as managing transactions in the sqlite3 module. Here are some quick references.
