如何利用上下文管理器扩展Python计时器

一个 Python 定时器上下文管理器

Python 有一个独特的构造，用于在代码块之前和之后调用函数：上下文管理器。

了解 Python 中的上下文管理器

上下文管理器长期以来一直是 Python 中重要的一部分。由 PEP 343 于 2005 年引入，并首次在 Python 2.5 中实现。可以使用 with 关键字识别代码中的上下文管理器：

with EXPRESSION as VARIABLE:
    BLOCK

EXPRESSION 是一些返回上下文管理器的 Python 表达式。首先上下文管理器绑定到变量名 VARIABLE上，BLOCK 可以是任何常规的 Python 代码块。上下文管理器保证程序在 BLOCK 之前调用一些代码，在 BLOCK 执行之后调用一些其他代码。这样即使 BLOCK 引发异常，后者也是会照样执行。

上下文管理器最常见的用途是处理不同的资源，如文件、锁和数据库连接等。上下文管理器用于使用资源后释放和清理资源。以下示例仅通过打印包含冒号的行来演示 timer.py 的基本结构。此外，它展示了在 Python 中打开文件的常用习语：

with open("timer.py") as fp:
    print("".join(ln for ln in fp if ":" in ln))

class TimerError(Exception):
class Timer:
    timers: ClassVar[Dict[str, float]] = {}
    name: Optional[str] = None
    text: str = "Elapsed time: {:0.4f} seconds"
    logger: Optional[Callable[[str], None]] = print
    _start_time: Optional[float] = field(default=None, init=False, repr=False)
    def __post_init__(self) -> None:
        if self.name is not None:
    def start(self) -> None:
        if self._start_time is not None:
    def stop(self) -> float:
        if self._start_time is None:
        if self.logger:
        if self.name:

注意，使用 open() 作为上下文管理器，文件指针fp 不会显式关闭，可以确认 fp 已自动关闭：

fp.closed

True

在此示例中，open("timer.py") 是一个返回上下文管理器的表达式。该上下文管理器绑定到名称 fp。上下文管理器在 print() 执行期间有效。这个单行代码块在 fp 的上下文中执行。

fp 是上下文管理器是什么意思？ 从技术上讲，就是 fp 实现了 上下文管理器协议。Python 语言底层有许多不同的协议。可以将协议视为说明我们代码必须实现哪些特定方法的合同。

上下文管理器协议由两种方法组成：

• 进入与上下文管理器相关的上下文时调用 .__enter__()。

• 退出与上下文管理器相关的上下文时调用 .__exit__()。

换句话说，要自己创建上下文管理器，需要编写一个实现 .__enter__() 和 .__exit__() 的类。试试 Hello, World!上下文管理器示例：

# studio.py
class Studio:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f"你好 {self.name}")
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        print(f"一会儿见, {self.name}")

Studio是一个上下文管理器，它实现了上下文管理器协议，使用如下：

from studio import Studio
with Studio("云朵君"):
    print("正在忙 ...")

你好 云朵君
正在忙 ...
一会儿见, 云朵君

首先，注意 .__enter__() 在做事之前是如何被调用的，而 .__exit__() 是在做事之后被调用的。该示例中，没有引用上下文管理器，因此不需要使用 as 为上下文管理器命名。

接下来，注意 self.__enter__() 的返回值受 as 约束。创建上下文管理器时，通常希望从 .__enter__() 返回 self 。可以按如下方式使用该返回值：

from greeter import Greeter
with Greeter("云朵君") as grt:
  print(f"{grt.name} 正在忙 ...")

你好 云朵君
云朵君 正在忙 ...
一会儿见, 云朵君

在写 __exit__ 函数时，需要注意的事，它必须要有这三个参数：

• exc_type：异常类型

• exc_val：异常值

• exc_tb：异常的错误栈信息

这三个参数用于上下文管理器中的错误处理，它们以 sys.exc_info() 的返回值返回。当主逻辑代码没有报异常时，这三个参数将都为None。

如果在执行块时发生异常，那么代码将使用异常类型、异常实例和回溯对象（即exc_type、exc_value和exc_tb）调用 .__exit__() 。通常情况下，这些在上下文管理器中会被忽略，而在引发异常之前调用 .__exit__()：

from greeter import Greeter
with Greeter("云朵君") as grt:
    print(f"{grt.age} does not exist")

你好 云朵君
一会儿见, 云朵君
Traceback (most recent call last):
  File "", line 2, in
AttributeError: 'Greeter' object has no attribute 'age'

可以看到，即使代码中有错误，还是照样打印了 "一会儿见, 云朵君"。

理解并使用 contextlib

现在我们初步了解了上下文管理器是什么以及如何创建自己的上下文管理器。在上面的例子中，我们只是为了构建一个上下文管理器，却写了一个类。如果只是要实现一个简单的功能，写一个类未免有点过于繁杂。这时候，我们就想，如果只写一个函数就可以实现上下文管理器就好了。

这个点Python早就想到了。它给我们提供了一个装饰器，你只要按照它的代码协议来实现函数内容，就可以将这个函数对象变成一个上下文管理器。

我们按照 contextlib 的协议来自己实现一个上下文管理器，为了更加直观我们换个用例，创建一个我们常用且熟悉的打开文件（with open）的上下文管理器。

import contextlib

@contextlib.contextmanager
def open_func(file_name):
    # __enter__方法
    print('open file:', file_name, 'in __enter__')
    file_handler = open(file_name, 'r')
 
    # 【重点】：yield
    yield file_handler

    # __exit__方法
    print('close file:', file_name, 'in __exit__')
    file_handler.close()
    return

with open_func('test.txt') as file_in:
    for line in file_in:
        print(line)

在被装饰函数里，必须是一个生成器（带有yield），而 yield 之前的代码，就相当于__enter__里的内容。yield 之后的代码，就相当于__exit__ 里的内容。

上面这段代码只能实现上下文管理器的第一个目的（管理资源），并不能实现第二个目的（处理异常）。

如果要处理异常，可以改成下面这个样子。

import contextlib

@contextlib.contextmanager
def open_func(file_name):
    # __enter__方法
    print('open file:', file_name, 'in __enter__')
    file_handler = open(file_name, 'r')

    try:
        yield file_handler
    except Exception as exc:
        # deal with exception
        print('the exception was thrown')
    finally:
        print('close file:', file_name, 'in __exit__')
        file_handler.close()
        return

with open_func('test.txt') as file_in:
    for line in file_in:
        1/0
        print(line)

Python 标准库中的 contextlib包括定义新上下文管理器的便捷方法，以及可用于关闭对象、抑制错误甚至什么都不做的现成上下文管理器！

创建 Python 计时器上下文管理器

了解了上下文管理器的一般工作方式后，要想知道它们是如何帮助处理时序代码呢？假设如果可以在代码块之前和之后运行某些函数，那么就可以简化 Python 计时器的工作方式。其实，上下文管理器可以自动为计时时显式调用 .start() 和.stop()。

同样，要让 Timer 作为上下文管理器工作，它需要遵守上下文管理器协议，换句话说，它必须实现 .__enter__() 和 .__exit__() 方法来启动和停止 Python 计时器。从目前的代码中可以看出，所有必要的功能其实都已经可用，因此只需将以下方法添加到之前编写的的 Timer 类中即可：

# timer.py
@dataclass
class Timer:
    # 其他代码保持不变

    def __enter__(self):
        """Start a new timer as a context manager"""
        self.start()
        return self

    def __exit__(self, *exc_info):
        """Stop the context manager timer"""
        self.stop()

Timer 现在就是一个上下文管理器。实现的重要部分是在进入上下文时， .__enter__() 调用 .start() 启动 Python 计时器，而在代码离开上下文时， .__exit__() 使用 .stop() 停止 Python 计时器。

from timer import Timer
import time
with Timer():
    time.sleep(0.7)

Elapsed time: 0.7012 seconds

此处注意两个更微妙的细节：

• .__enter__() 返回 self，Timer 实例，它允许用户使用 as 将 Timer 实例绑定到变量。例如，使用 with Timer() as t: 将创建指向 Timer 对象的变量 t。

• .__exit__() 需要三个参数，其中包含有关上下文执行期间发生的任何异常的信息。代码中，这些参数被打包到一个名为 exc_info 的元组中，然后被忽略，此时 Timer 不会尝试任何异常处理。

在这种情况下不会处理任何异常。上下文管理器的一大特点是，无论上下文如何退出，都会确保调用.__exit__()。在以下示例中，创建除零公式模拟异常查看代码功能：

from timer import Timer
with Timer():
    for num in range(-3, 3):
        print(f"1 / {num} = {1 / num:.3f}")

1 / -3 = -0.333
1 / -2 = -0.500
1 / -1 = -1.000
Elapsed time: 0.0001 seconds
Traceback (most recent call last):
  File "", line 3, in
ZeroDivisionError: division by zero

注意 ，即使代码抛出异常，Timer 也会打印出经过的时间。

使用 Python 定时器上下文管理器

现在我们将一起学习如何使用 Timer 上下文管理器来计时 "下载数据" 程序。回想一下之前是如何使用 Timer 的：

# download_data.py
import requests
from timer import Timer
def main():
    t = Timer()
    t.start()
    source_url = 'https://cloud.tsinghua.edu.cn/d/e1ccfff39ad541908bae/files/?p=%2Fall_six_datasets.zip&dl=1'
    headers = {'User-Agent': 'Mozilla/5.0'}
    res = requests.get(source_url, headers=headers) 
    t.stop()
    with open('dataset/datasets.zip', 'wb') as f:
        f.write(res.content)

if __name__ == "__main__":
    main()

我们正在对 requests.get() 的调用进行记时监控。使用上下文管理器可以使代码更短、更简单、更易读：

# download_data.py
import requests
from timer import Timer
def main():
    source_url = 'https://cloud.tsinghua.edu.cn/d/e1ccfff39ad541908bae/files/?p=%2Fall_six_datasets.zip&dl=1'
    headers = {'User-Agent': 'Mozilla/5.0'}
    with Timer():
        res = requests.get(source_url, headers=headers)
        
    with open('dataset/datasets.zip', 'wb') as f:
        f.write(res.content)

if __name__ == "__main__":
    main()