Python 有很多魔法方法,本文记录一下可以自定义 with 语句的上下文管理器所使用到的两个魔法方法,也就是 __enter____exit__ 方法的实用性。

自定义上下文管理类

最常见的 with 语句就是 open 函数了,这里不做解释,直接来看一个自定义类的例子。

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)

上面定义了一个类,这个类定义了两个打印值的方法,其中一个方法会报错,同时该类使用到了两个魔法方法,有了这两个方法,这个类就可以使用 with 语句来进行调用,来看看调用正常函数的结果:

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

结果如下:

2
exc_type: None
exc_val: None
exc_tb: None

再来看看调用报错函数的结果

with TestHandler() as t:
    t.bad_func()
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

从上面两次调用,可以看到,__exit__ 函数里面的三个参数(定义函数的时候默认会要求加入)分别代表了报错类型、报错原因、报错追溯,只有当 with 语句调用报错时候,这三个参数才有值,否则就是 None,看到这里,你是否能够想到什么?可以利用这三个参数进行异常判断和处理。

上下文管理实用性

已经知道如何定义 with 语句了,也知道遇到异常会出现什么,那么现在来看看自定义 with 语句的使用场景有哪些。

with 语句比较适合的场景是打开->操作->关闭,在我们常用的除了文件操作外,还有数据库操作、SSH 操作会涉及这个过程,所以,直接看看这两个操作的例子。

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()

上面这个关于数据操作的例子就很典型,它包括了数据库连接、数据库操作(with 语句之后)、异常处理、数据库关闭连接等操作。

看一下 with 语句的使用:

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

是不是非常的方便,当然,如果再结合 try 语句来进行连接操作,就更安全可靠。

再来看看 SSH 的操作例子

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()

很明显,上面的自定义类 with 语句返回的是一个 SSHClient 对象,所以使用时直接按照这个对象的方法调用即可,调用结束会自动断开连接。

总结:自定义 with 语句简单理解就是非常适合一些“有始有终”的场景,通过自定义上下文管理器,可以把一些需要重复执行的固定操作简化,只需要关注特定的操作本身。