The world’s leading publication for data science, AI, and ML professionals.

What Is Behind The Python “with” Statement?

You can write customized function or classes supporting the "with" statement

Image by Rudy and Peter Skitterians from Pixabay
Image by Rudy and Peter Skitterians from Pixabay

I’m sure that you must ever use the "with" statement in Python, as long as you are not a Python newbie who just started to learn it a few days ago. It is commonly known that we can use the with-statement to create a "context" of a resource without being worried to close or release it afterwards.

Since this syntax is kind of unique in Python, have you ever think what exactly happened behind the scene? In this article, I’ll introduce Python with statements in-depth and provide some ideas to flexibly use it in your program.

1. Why Use the "with" Statement?

Image by Juraj Varga from Pixabay
Image by Juraj Varga from Pixabay

Let’s start with something basic. One of the most common use cases of the with-statement is opening a file in the Python problem. Suppose we want to open a text file and write something into it, and then close it. The simplest code is as follows.

1.1 A Bad Example

# Bad Example
f = open('my.txt', 'w')
f.write('Hello!')
f.close()

The above code will open a file called my.txt and then write the string Hello! into it. Then, we need to close the file stream to release this resource to avoid issues such as memory leaks.

However, this is a bad example. If there are something goes wrong during we manipulate the file, the program will never reach the close() method.

1.2 A "Regular" Solution

So, a better solution is to put the close() method in the finally-block followed by a try-except block.

# Try Except
f = open('my.txt', 'r')
try:
    print(f.readline())
except:
    pass
finally:
    f.close()

The above code opens the file in read-mode this time, then it read the first line and prints it. The "finally" block ensures that the close() method is guaranteed to be executed even if there could be errors.

1.3 The Recommended Solution

The recommended solution is to use the with-statement in Python. The code is as simple as follows.

# With Statement
with open('my.txt', 'r') as f:
    print(f.readline())

We opened the file as the variable f, then we can use this file inside this with-statement block to do whatever we want. We won’t be bothered to close the file. It is also guaranteed to be destroyed eventually. Also, we cannot use the variable f outside the context.

It is also worth mentioning that we can put multiple comma-separated context expressions in a with-statement since Python 3.1.

with open('my.txt', 'r') as f_source, open('my_new.txt', 'w') as f_dest:
    f_dest.write(f_source.readline() + ' Again')

Also, since Python 3.10, we can group multiple expressions using parentheses, so that the code can be broken up into multiple lines and improve the readability.

with (
    open('my.txt', 'r') as f_source,
    open('my_new.txt', 'w') as f_dest
):
    f_dest.write(f_source.readline() + ' Again')

By the way, the above code opens the original file my.txt as the source and a new file my_new.txt as the destination. Then, it reads from the source file and writes to the new file as the destination.

2. "With" Customized Classes

Image by bodobe from Pixabay
Image by bodobe from Pixabay

All the retrospections have been done in the previous section. Now we should begin with something in-depth. Why does the open() method support with-statement? Can we create our customised classes that support it too?

2.1 A Simple Example

The answer is yes. Let’s start with a simple class as an example.

class Person:
    def __init__(self, name):
        self.name = name

    def greeting(self):
        print(f'Hello, my name is {self.name}!')

The class Person has only two methods, one is the init method that takes the name value as a property of the class. Then, there is also a greeting() method that prints a greeting message using the name of this person.

Let’s try its greeting() method.

Person('Chris').greeting()

Now, suppose that we want to create this person in a context using the with-statement. Can we do that?

with Person('Chris') as p:
    p.greeting()

Apparently, we can’t. However, the error reminded us what should we do.

2.2. Implementation of the enter and exit Methods

To be able to use the with-statement, we need to implement the __enter__() method. In fact, there is also another method called __exit__() that is optional. Let’s see the code below.

class Person:
    def __init__(self, name):
        self.name = name

    def greeting(self):
        print(f'Hello, my name is {self.name}!')
    def __enter__(self):
        print('Creating a person object in a context...')
        return self
    def __exit__(self, exc_type, exc_value, exc_tb):
        print("exception type:", exc_type)
        print("exception value:", exc_value)
        print("exception traceback:", exc_tb)
        print('Object disposed')

Here, the __enter__() method will be executed when the with-statement is used to create the object. It must return the object, which is itself.

The __exit__() method will be called when the context is finished and the object needs to be disposed. It also has 3 parameters that will catch some details of exceptions if there are any.

Now, let’s try the code again.

with Person('Chris') as p:
    p.greeting()

From the output, we can see that the __enter__() method is called before anything else, and then the body of the with statement is executed, which is the greeting message. Finally, the __exit__() method is called and the object is disposed.

2.3 Handling Exceptions in the With-Statement

There was no exception happened in the code above, so all of the exception-related parameters are None.

with Person('Chris') as p:
    a = 1/0
    p.greeting()

In this piece of code, I have added a = 1/0 which will cause errors for sure. Let’s see how with-statement will behave.

The exception type, value and traceback are available now because there was an exception. These callback parameters enable us to handle the exceptions accordingly.

3. "With" Customised Functions

Image by Uwe Baumann from Pixabay
Image by Uwe Baumann from Pixabay

Well, most of the time we may want to use the with-statement with a function rather than creating an object. The open() function is such an example. That is, we write a with-statement with this function and give the filename that we want to open. Then, it opens the file within this context and the connection is destroyed when leaving the context.

We can do this using the contextlib that is built-in to Python. Specifically, we need to use the contextmanager in this module.

from contextlib import contextmanager

Then, we can define the function with the contextmanager as decorator.

Suppose we are going to establish a connection with a database. Let’s do this using pseudo code.

@contextmanager
def get_db_connection(connection_str):
    db_con = 'connected to' + connection_str
    try:
        yield db_con
    finally:
        print('the connection is closed')

The code simulated connecting to a database. Please be noticed that we need to use yield keyword to return the database connection. Also, we need to use the try-except block and make sure that the exceptions are handled, and the resource should be released in the finally block.

Now we can use this function in a with-statement.

with get_db_connection('mysql') as con:
    print(con)

Summary

Image by garageband from Pixabay
Image by garageband from Pixabay

In this article, I have introduced when we need to use a with-statement and how to use it. The mechanism behind this syntax was also explained. Apart from that, I have also introduced how to write a class or a function that supports the with-statement. By implementing a class or a function that way, we can simplify the use of the interface in a very Pythonic way.

Join Medium with my referral link – Christopher Tao

If you feel my articles are helpful, please consider joining Medium Membership to support me and thousands of other writers! (Click the link above)

Unless otherwise noted all images are by the author


Related Articles