with Statement in Python

The with statement in Python helps you with resource management. It ensures no resources are accidentally left open.

The with statement is a replacement for commonly used try/finally error-handling statements.

A common example of using the with statement is opening a file.

To open and write to a file in Python, you can use the with statement as follows:

with open("example.txt", "w") as file:
    file.write("Hello World!")

The with statement ensures the file is closed after the writing has been completed.

Another way to write to a file is like this:

f = open("example.txt", "w")

try:
    f.write("hello world")
finally:
    f.close()

But as you can see, the with approach is way smoother.

As it turns out, the latter resembles what happens behind the scenes when you call open() using with statement.

Today you are going to:

Learn what is the with statement.
Learn what are context managers.
Implement a context manager class and a context manager method.
Finally, you will get creative with context managers.

How to Open Files in Python

To open and write to a file, you can use this approach:

f = open("example.txt", "w")

try:
    f.write("hello world")
finally:
    f.close()

In this piece of code:

The file is opened separately using the open() function.
Exception handling is used to write "hello world" is into the opened file.
The file is manually closed in the finally block.

This small program opens up a file called example.txt (assuming it is in the same folder as the project file). If there is no such file, a new file is created. Then the program writes "hello world" to the file and closes it.

There is nothing wrong with this approach. But there is a more elegant way to do it—by using the with statement.

Let’s repeat the above example using the with statement:

with open("example.txt", "w") as file:
    file.write("hello world")

This simplifies the code—you can let the with statement take care of closing the file after being used. This is what makes using the with statement a recommended way to open files in general.

But what does the with statement do? Can you call it on any object?

Let’s next discuss the with statement in more detail.

Using the with Statement in a Custom Class—Context Managers in Python

One might think the with statement only works with the open() function. But this is not the case.

You can create your own classes and objects that support the with statement. A class or function that supports the with statement is known as a context manager.

You may want to write a context manager if you want to automatically close resources in your project.

For a class to qualify as a context manager, it needs to implement these two methods:

__enter__()
__exit__()

After implementing these methods, you can use the with statement on the objects of the class.

When you call the with statement, the __enter__() method is invoked.
When you exit the scope of the with block, the __exit__() function is called.

For example, let’s create a file writer context manager. This class works similar to the open() method:

class FileWriter(object):
    def __init__(self, file_name):
        self.file_name = file_name

    def __enter__(self):
        self.file = open(self.file_name, "w")
        return self.file

    def __exit__(self, exception_type, exception_value, traceback):
        self.file.close()

Now you can use this custom class with the with statement as follows:

with FileWriter("example.txt") as file:
    file.write("hello world")

As a result, this program creates a file called example.txt and writes "hello world" into it—Just like the open() function.

Let’s go through how the code works.

with FileWriter("example.txt") creates a new FileWriter object and calls __enter__().
The __enter__() method initializes the resource you want to use. In this case, it opens a text file. It also has to return the descriptor of the resource, so it returns the opened file.
The as file assigns the file to a variable file.
Finally, the code you want to run with the acquired resource is placed in the with block after the colon.
As soon as this code finishes execution, the __exit__() method is automatically called. In this case, it closes the file.

Write Your Own Context Manager Methods in Python

So far you have learned:

What is a context manager in Python.
How to write a context manager class.
How to use the with statement on a context manager object.

The context manager you previously wrote is a class. What if you want to create a context manager method similar to the open() instead?

This is simple. Python supports writing custom context manager methods.

To turn a method into a context manager, use the contexlib module:

Import the contexlib module
Mark a method as a context manager with a contexlib.contextmanager decorator.

As an example, let’s create a custom version of the open() method. This function opens a file, writes text to it, and closes it.

from contextlib import contextmanager

@contextmanager
def my_open(name):
    try:
        file = open(name, "w")
        yield file
    finally:
        file.close()

with my_open("example.txt") as file:
    file.write("hello world")

This code works such that:

The my_open() function is marked as a context manager using contextmanager decorator.
The function opens a file.
Then it pauses the execution and hands the opened file over to the caller using yield. (If return was used, the file would close immediately. Learn more about using yield in Python.)
When code inside thewith block completes, the my_open() function continues execution and closes the file by running the finally block.

So far, you have learned:

What is a context manager.
How to use context managers using the with statement.
How to write a context manager class.
How to write a context manager function.

But in the previous examples, you’ve only dealt with files. This is not the only use case for context managers. You can also get creative.

Let’s see an example.

Context Manager Example—Get Creative

A context manager is a class or method that supports the with statement. In other words, it implements the methods __enter__() and __exit__().

Context managers are flexible. You can get creative with them. You are not restricted to using context managers with files only.

Example

Let’s create a bullet point list object that supports any number of indentations.

Obviously, there are many ways to achieve this. But as we are talking about context managers, let’s write one for this task.

The goal is to create a program that supports a nested with structure that acts as a nested bulleted list. This would make the code of writing bulleted lists readable:

with BulletedList() as bp:
    bp.item("Dessert")
    with bp:
        bp.item("Apple Pie")
        with bp:
            bp.item("Apples")
            bp.item("Cinamon")
            bp.item("Sugar")
            bp.item("Flour")
            bp.item("Eggs")
    with bp:
        bp.item("Hot Chocolate")
        with bp:
            bp.item("Milk")
            bp.item("Chocolate powder")
            bp.item("Cream")

As a result, we want an output like this:

  - Dessert
    * Apple Pie
      o Apples
      o Cinamon
      o Sugar
      o Flour
      o Eggs
    * Hot Chocolate
      o Milk
      o Chocolate powder
      o Cream

Let’s turn this idea into code.

Implementation

To write a program that behaves as described, let’s write a context manager.

The context manager takes care of the indentations of the bulleted list elements. It also chooses the right bullet type.

Here is how to put it together in code:

class BulletedList:
    def __init__(self):
        self.indent_level = 0

    def __enter__(self):
        self.indent_level += 1
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.indent_level -= 1

    def bullet_type(self):
        bullet_types = ["o ", "- ", "* "]
        return bullet_types[self.indent_level % 3]

    def item(self, text):
        print("  " * self.indent_level + self.bullet_type() + text)

Now you can use this class to create bulleted lists according to the specifications.

Let’s next inspect how the code works:

The BulletedList is a context manager for creating bulleted lists of multiple indentations.
A BulletedList object is initialized with the indentation level of 0.
When a BulletedList is used with a with statement, the __enter__() function is called under the hood. This increases the indentation level by one.
The item() function prints the bulleted list item into the correct indentation level. The bullet type is given by the bullet_type() function.
When a with block finishes running, the indentation level decreases as the __exit__() gets called automatically.

Notice how this task had nothing to do with opening/closing files. Instead, the __enter__() and __exit__() methods took care of the indentation levels.

Conclusion

In Python, the with statement replaces a try-catch block with a concise shorthand. More importantly, it ensures closing resources right after processing them.

A common example of using the with statement is reading/writing to a file.

A function or class that supports the with statement is known as a context manager.

A context manager allows you open and close resources right when you want to.

For example, the open() function is a context manager. When you call the open() function using the with statement, the file closes automatically after you have processed the file.

You may also write your own context manager classes and methods.

A context manager class needs to implement __enter__() and __exit__() methods.
A context manager function is implemented using the contexlib module’s contextmanager decorator.

Feel free to get creative with context managers—you are not restricted to only dealing with files with them.

Thanks for reading. Happy coding!

The with Statement in Python