Python has a lot of magic methods, and this article documents the usefulness of two magic methods used by context managers that can be customized with statements, namely the __enter__ and __exit__ methods.

Custom Context Management Classes

The most common with statement is the open function. Instead of explaining it here, let’s look directly at an example of a custom class.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

class TestHandler():
    def __init__(self):
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('exc_type:', exc_type)
        print('exc_val:', exc_val)
        print('exc_tb:', exc_tb)

    def func(self):
        print(1 + 1)

    def bad_func(self):
        print('a' + 1)

The above defines a class that defines two methods that print values, one of which throws an exception. Also the class uses two magic methods. With these two methods, the class can be called using the with statement to see the result of calling a normal function.

1
2

with TestHandler() as t:
    t.func()

The results are as follows.

1
2
3
4

2
exc_type: None
exc_val: None
exc_tb: None

Then look at the result of calling a function that throws an exception.

1
2

with TestHandler() as t:
    t.bad_func()

1
2
3
4
5
6
7
8
9

exc_type: <class 'TypeError'>
exc_val: Can't convert 'int' object to str implicitly
exc_tb: <traceback object at 0x0000021CEB484B08>
Traceback (most recent call last):
  File "D:/Mycode/TestCase/mark.py", line 23, in <module>
    t.bad_func()
  File "D:/Mycode/TestCase/mark.py", line 17, in bad_func
    print('a' + 1)
TypeError: Can't convert 'int' object to str implicitly

From the above two calls, you can see that the three parameters inside the __exit__ function (which are required by default when defining the function) represent the type of error, the cause of the error, and the trace of the error, and only when the with statement is called with an error, these three parameters have values, otherwise they are None. These three parameters can be used to determine and handle exceptions.

Context Management Utilities

We already know how to define the with statement and what happens when we encounter an exception, so now let’s see what are the scenarios for using the custom with statement.

The with statement is more suitable for the scenario of open->operate->close, in addition to our common file operations, there are database operations, SSH operations will involve this process. So, let’s take a look at the examples of these two operations directly.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

import sqlite3

class DBHandler():

    def __init__(self, database):
        self.database = database
        self.conn = sqlite3.connect(self.database)
        self.cursor = self.conn.cursor()

    def __enter__(self):
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.conn.commit()
        self.conn.close()

The above example of data manipulation is typical, it includes database connection, database manipulation (after the with statement), exception handling, and database connection closure.

Take a look at the use of the with statement.

1
2

with DBHandler(database) as db:
    db.executescript(create_sql)

It is very convenient, but of course, if you combine it with try statements to connect, it will be more secure and reliable.

Let’s take a look at the SSH operation example.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

import paramiko

class SSHHandler():
    def __init__(self, host, username, password, port=22):
        self.host = host
        self.username = username
        self.password = password
        self.port = port
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy)
        self.ssh.connect(self.host, self.port, self.username, self.password)

    def __enter__(self):
        return self.ssh

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.ssh.close()

Obviously, the above custom class with statement returns an SSHClient object, so it can be called directly according to the methods of this object, and will be automatically disconnected at the end of the call.

To summarize: the custom with statement is simply understood to be very suitable for some “start and finish” scenarios, by customizing the context manager, you can simplify some fixed operations that need to be executed repeatedly, and only need to focus on the specific operation itself.