深入理解Python上下文管理-with语句

深入理解Python上下文管理-with语句

由于最近一直在阅读Python OS模块的源码,学到了很多优雅的Python语法使用。对于Python中的上下文管理很早之前就了解过。但是一直没有深入理解过,趁着学习标准库源码之际,好好研究一下。

对于Python中with的用法,我觉得很多人一开始都是从with open( file ) as … 这个语句开始接触到的,包括我自己也是。刚开始学习Python的时候,都知道这是一种非常优雅的写法,因为有一些任务,可能事先需要设置,事后做清理工作。对于这种场景,Python的with语句提供了一种非常方便的处理方式。

这是os模块中walk函数的部分代码,通过scandir()是nt模块中的遍历当前目录文件的一个函数,os.walk()算是对它的一个更完善的封装。这里只是抛砖引玉,标准库的代码非常优秀。


Python2.5之后引入了上下文管理器(context manager),算是Python的黑魔法之一,它用于规定某个对象的使用范围。

为什么需要上下文管理器?

正常情况下,管理各种系统资源(如文件),数据库连接时,通常是先打开这些资源,执行完相应的业务逻辑,最后关闭资源。

举两个例子:

  1. 使用Python打开一个文件写入内容,之后需要关闭这个文件。如果不正常关闭的话可能会在文件操作时出现异常,因为系统允许你打开的文件的最大数是有限的。
  2. 在数据库连接时也是存在类似问题,数据库的连接算是一种比较昂贵的资源,若连接过多而没有及时关闭的话,就可能出现不能继续连接的异常错误。

但是,很多程序员经常会忘记关闭文件,或者关闭数据库的连接。这时候就引入了上下文管理器,它可以在你不需要该对象的时候,自动关闭它。这里我引用了一段别人的解释帮助理解。


文件处理就是一个很好的例子,需要先获取一个文件句柄,从文件中读取数据,然后关闭文件句柄。

file  = open('zgao.py')
data = file.read()
file.close()

这算是最原始的操作文件的写法,但是每次都要写一个close的操作关闭句柄。非常容易忘掉,另外文件读取数据发生异常,没有进行任何处理。

try:
    f = open('zgao.py')
except:
    print('fail to open')
    exit(-1)
finally:
     f.close()

这是添加了异常处理之后的写法,虽然更加完善了,但是代码看起来很冗余,不符合Python一贯简洁的风格。使用with,除了有更优雅的语法,with还可以很好的处理上下文环境产生的异常。

with的工作原理

  • 紧跟with后面的语句被求值后,返回对象的 __enter__() 方法被调用,这个方法的返回值将被赋值给as后面的变量。
  • 当with后面的代码块全部被执行完之后,将调用前面返回对象的 __exit__()方法。

自定义的上下文管理器要实现上下文管理协议所需要的 __enter__() 和 __exit__() 两个方法:

  • context_manager.__enter__() :进入上下文管理器的运行时上下文,在语句体执行前调用。with 语句将该方法的返回值赋值给 as 子句中的 target,如果指定了 as 子句的话
  • context_manager.__exit__(exc_type, exc_value, exc_traceback) :退出与上下文管理器相关的运行时上下文,返回一个布尔值表示是否对发生的异常进行处理。参数表示引起退出操作的异常,如果退出时没有发生异常,则3个参数都为None。如果发生异常,返回True 表示不处理异常,否则会在退出该方法后重新抛出异常以由 with 语句之外的代码逻辑进行处理。如果该方法内部产生异常,则会取代由 statement-body 中语句产生的异常。要处理异常时,不要显示重新抛出异常,即不能重新抛出通过参数传递进来的异常,只需要将返回值设置为 False 就可以了。之后,上下文管理代码会检测是否 __exit__() 失败来处理异常。
class Zgao:
    def __init__(self):
        print('对象初始化')
    def __enter__(self):
        print("正在执行 __enter__()")
        return "zgao"
    def __exit__(self, type, value, trace):
        print("正在执行 __exit__()")
 
context = Zgao()
 
with context as f:
    print("__enter__中的返回值赋值给了as后的f:", f)

我单独写的一个类运行后来解释原理。

我先实例化了一个对象context,通过with语句将返回值赋给f。从打印的顺序就可以理解with语句的执行顺序了。但是with真正强大之处是它可以处理异常。可能你已经注意到Zgao这个类的 __exit__ 方法有三个参数 val, type 和 trace。 这些参数在异常处理中相当有用。我们稍微修改一下代码。

在with后面的代码块抛出任何异常时,__exit__() 方法被执行。异常抛出时,与之关联的type,value和stack trace传给 __exit__() 方法,因此抛出的ZeroDivisionError异常被打印出来了。所以清理资源,关闭文件等等操作,都可以放在 __exit__ 方法当中。

另外,__exit__ 除了用于tear things down,还可以进行异常的监控和处理,注意后几个参数。要跳过一个异常,只需要返回该函数True即可。

下面的修改一下代码代码,遇到ZeroDivisionError就忽略掉,而让其他异常正常抛出。其中isinstance就是用于判断两种是否为同类异常,如果第二个参数是他的父类,也返回True。

__exit__ 函数可以进行部分异常的处理,如果我们不在这个函数中处理异常,他会正常抛出。

通过上下文管理器,可以更好的控制对象在不同区间的特性,并且可以使用with语句替代try…except方法,使得代码更加的简洁,主要的使用场景是访问资源,可以保证不管过程中是否发生错误或者异常都会执行相应的清理操作,释放出访问的资源。

with-as表达式极大的简化了每次写finally的工作,这对保持代码的优雅性是有极大帮助的。

但是到这里就结束了吗?不,你会不会觉得每次要用with语句都需要在类里面加入__enter__和__exit__,这样非常麻烦呢。是的,采用传统方式创建上下文管理器,即编写一个包含 __enter__() 和 __exit__() 方法的类,这并不难。不过有些时候,对于很少的上下文来说,完全编写所有代码会是额外的负担。在这些情况下,可以使用 contextmanager() 修饰符将一个生成器函数转换为上下文管理器。

contextlib 模块

contextlib 模块提供了3个对象:装饰器 contextmanager、函数 nested 和上下文管理器 closing。使用这些对象,可以对已有的生成器函数或者对象进行包装,加入对上下文管理协议的支持,避免了专门编写上下文管理器来支持 with 语句。

装饰器 contextmanager

contextmanager 用于对生成器函数进行装饰,生成器函数被装饰以后,返回的是一个上下文管理器,其 __enter__() 和 __exit__() 方法由 contextmanager 负责提供,而不再是之前的迭代子。被装饰的生成器函数只能产生一个值,否则会导致异常 RuntimeError;产生的值会赋值给 as 子句中的 target,如果使用了 as 子句的话。

可以看到,生成器函数中 yield 之前的语句在 __enter__() 方法中执行,yield 之后的语句在 __exit__() 中执行,而 yield 产生的值赋给了 as 子句中的 result变量。

需要注意的是,contextmanager 只是省略了 __enter__() / __exit__() 的编写,但并不负责实现资源的“获取”和“清理”工作;“获取”操作需要定义在 yield 语句之前,“清理”操作需要定义 yield 语句之后,这样 with 语句在执行 __enter__() / __exit__() 方法时会执行这些语句以获取/释放资源,即生成器函数中需要实现必要的逻辑控制,包括资源访问出现错误时抛出适当的异常。

赞赏

微信赞赏支付宝赞赏

Zgao

愿有一日,安全圈的师傅们都能用上Zgao写的工具。

发表评论