Python 上下文管理器编程模式写出优雅代码——@contextmanager 装饰器

作者：高玉涵

时间：2023.11.3 22:34

环境：Python 3.10.4

既然我们改变不了世界，那么就努力不让世界改变。——《甜蜜家园》

Python 布道者

​ Python 语言中有一些不常见的特性，正因如此 Python 用户往往会忽视或没有充分使用这些特性。with 语句和上下文管理器即是其中之一（在各种语言中 with 语句的作用不同，不要觉得名字一样，就意味着作用也一样），with 语句会设置一个临时的上下文。如果你比较健忘，或是你的程序有太多元素需要管理，如”打开之后需要关闭“、”锁上之后需要释放“或是”修改之后需要还原“的场景。交给上下文管理器对象控制并负责清理上下文，这么做能避免错误并减少样板代码。

现实中的我

​ 初看上面这段话时心无波澜，只到近期因工作需要处理大量数据，遂编写程序完美解决之。得暇，回望先前写的代码，方才发现其中充斥着大量的样板代码。霎时脑海里浮现了开篇的那段话。

举个现实中的例子

​ 日常编写程序，养成了将程序运行状态实时记录到日志的习惯。如，模块、方法执行状态及耗时情况，以便定位结果与预期不符发生的地方，某方法较耗时时，做针对性优化等。这些记录状态的语句，通常格式固定且遍布在代码的各个角落，也就是所谓的样板代码。举例代码如下：

class Card:
    def __init__(self):
        self.__start = default_timer() # 开始时间
		
        # 此间,代码省略...
        
        diff = default_timer() - self.__start
        self.__logger.info('%s', f'耗时 {diff:.6f} 秒')
	
    def read_file(self):
        self.__start = default_timer() # 开始时间
		
        # 此间,代码省略...
        
        diff = default_timer() - self.__start
        self.__logger.info('%s', f'耗时 {diff:.6f} 秒')
        
    def to_card(self):
        self.__start = default_timer() # 开始时间
		
        # 此间,代码省略...
        
        diff = default_timer() - self.__start
        self.__logger.info('%s', f'耗时 {diff:.6f} 秒')
        
    def statis_card(self):
        ......
        
# 此处,省略其它类方法

​ 这里假设你已知道并成功初始化 logging 对象。类中每个方法，都会记录执行所耗时间，并在结束时输出到日志。日志大概样式：

2023-11-03 21:46:37,428 - __init__ - INFO - 耗时 0.000489 秒
2023-11-03 21:46:37,685 - read_file - INFO - 去重 0 条记录.耗时 0.256535 秒
2023-11-03 21:47:13,865 - to_card - INFO - 耗时 0.113831 秒
2023-11-03 21:47:19,491 - statis_card - INFO - 耗时 2.108383 秒
......

contextlib 模块中的 @contextmanager 装饰器

​ 一般来说，要创建一个上下文管理器的话，需要创建一个带有 __enter__ 和 __exit__ 函数的类。 __enter__ 函数负责返回要管理的资源（例如文件或 socket），__exi__ 函数负责执行清理工作（例如关闭文件）。

​ 但是，如果你的应用场景不需要进行这么细致的管理，也可以使用 @contextlib.contextmanager 来创建简单的上下文管理器，用它把一个生成器函数转换为上下文管理器。修改后的代码：

import contextlib

class Card:
    # 此间,代码省略...
    
    @contextlib.contextmanager
    def __time_consuming(self, text:str):
        try:
            self.__start = default_timer()   # 开始时间
            yield
        finally:
            diff = default_timer() - self.__start
            self.__logger.info('%s', f'{text} 耗时 {diff:.6f} 秒')
   
# 此处,省略其它类方法

​ 留意一下 __time_consuming 函数定义中出现的 try 和 finally 代码块。你可能经常遇到 try/except 语句，但是 try/finally 却不那么常见。不管 try 中出现什么异常， finally 代码块最后一定会被执行。这里我们需要这个特性，不管代码如何运行，最后都将把结果输出到日志。这个装饰器也有迷惑人的一面，因为它与迭代无关，却要使用 yield 语句，这里你简单把它理解成 return 同样的功能即可。代码：

class Card:
    def __init__(self):
        with self.__time_consuming('__init__'):
        	# 此间,代码省略...
        
    def read_file(self):
        with self.__time_consuming('read_file'):
        	# 此间,代码省略...
        
    def to_card(self):
        with self.__time_consuming('to_card'):
        	# 此间,代码省略...
        
    def statis_card(self):
        ......
        
# 此处,省略其它类方法

​ @contextmanager 装饰器优雅且实用，把三个不同的 Python 特性结合到了一起：函数装饰器、生成器和 with 语句。