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: exc_val: Can't convert 'int' object to str implicitly exc_tb: Traceback (most recent call last): File "D:/Mycode/TestCase/mark.py", line 23, in 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.