两种方法自定义Python上下文管理器-contextlib

今天分享两种自定义上下文管理器的方法，并比较它们的性能。上下文管理器相信你也用过，就是 with 开头的 Python 代码块，通常在读写文件、数据库的时候必用，但是我自己很少自己写，主要还是没有意识到它的好，但是今天，我意识到了。比如说，编程时经常要处理的路径问题，通常程序工作在一个目录，但是过程中要去另一个目录处理一些文件，处理完还要切换回原目录，这就需要不停的调用 os.chdir：

import os

origin_path = os.getcwd()

os.chdir(new_path)
do_something()
os.chdir(origin_path)

...

os.chdir(new_path2)
do_something2()
os.chdir(origin_path)
...

如果切换多次，你会看到满天飞的 os.chdir，这一点也不 Pythonic。接下来分享两种自定义上下文管理器的方法，然后你就会知道上下文管理器（contextlib) 的优雅和便捷。

方法一、使用 contextlib.contextmanager 装饰器

from contextlib import contextmanager import os @contextmanager def change_path(new_path):     origin_path = os.getcwd()
    os.chdir(new_path)
    try:
        yield os.getcwd()
    finally:
        os.chdir(origin_path) if __name__ == "__main__":
    print(f"with 之前的目录 {os.getcwd()}")
    with change_path("/Users/aaron/tmp") as new_path:
        print(f"with 之中的目录 {os.getcwd()}")
        print("do something...")
    print(f"with 之后的目录 {os.getcwd()}")

上面的代码中 change_path 就是一个上下文管理器，在 mian 中，完全不需要对目录切换来切换去，运行结果如下：

with 之前的目录 /Users/aaron/gitee/somenzz
with 之中的目录 /Users/aaron/tmp do something...
with 之后的目录 /Users/aaron/gitee/somenzz

方法二、使用类来实现装饰器

import os class ChangePath:     def __init__(self, new_path) -> None:         self.new_path = new_path

    def __enter__(self):         self.origin_path = os.getcwd()
        os.chdir(self.new_path)
        return os.getcwd()
    
    def __exit__(self,exc_type,exc_value,traceback):         os.chdir(self.origin_path) if __name__ == "__main__":
    print(f"with 之前的目录 {os.getcwd()}")
    with ChangePath("/Users/aaron/tmp") as new_path:
        print(f"with 之中的目录 {os.getcwd()}")
        print("do something...")
    print(f"with 之后的目录 {os.getcwd()}")

可以看出，只要一个类实现了 __enter__ 和 __exit__ 方法，就可以当做上下文管理器。运行结果是一样的：

with 之前的目录 /Users/aaron/gitee/somenzz
with 之中的目录 /Users/aaron/tmp do something...
with 之后的目录 /Users/aaron/gitee/somenzz

哪种方法更快？

测试之前，我直觉认为，方法二会稍微快一点，因为我很清楚进入和离开 contextlib 时，程序都做了些什么。不过还是用事实说话，测试完整代码：

from contextlib import contextmanager import os from timeit import timeit @contextmanager def change_path(new_path):     origin_path = os.getcwd()
    os.chdir(new_path)
    try:
        yield os.getcwd()
    finally:
        os.chdir(origin_path) class ChangePath:     def __init__(self, new_path) -> None:         self.new_path = new_path

    def __enter__(self):         self.origin_path = os.getcwd()
        os.chdir(self.new_path)
        return os.getcwd()

    def __exit__(self, exc_type, exc_value, traceback):         os.chdir(self.origin_path) def method1():     os.getcwd()
    with change_path("/Users/aaron/tmp") as new_path:
        os.getcwd()
    os.getcwd() def method2():     os.getcwd()
    with ChangePath("/Users/aaron/tmp") as new_path:
        os.getcwd()
    os.getcwd() if __name__ == "__main__":
    number = 1000     x1 = timeit(
        "method1()", setup="from __main__ import method1", number=number
    )
    x2 = timeit(
        "method2()", setup="from __main__ import method2", number=number
    )
    print(f"方法一执行 {number} 次耗时 {x1 :.8f}")
    print(f"方法二执行 {number} 次耗时 {x2 :.8f}")

我执行了三次，分别测试了 1000，10000，100000 次，均是方法二更快：

最后的话

进入代码前需要做一些预处理，离开后需要再次处理的时候，上下文管理器就可以出场了，如果这样的处理变多的时候，你会发现 contextlib 真的是个好东西，大大减轻了程序员的心智和记忆负担。本文分享了两种实现 contextlib 的方法，推荐使用类来实现 contextlib。想要了解更多关于Python的教程欢迎持续关注编程学习网