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

somenzz

发布于 2022-05-24 14:43:36

3360

发布于 2022-05-24 14:43:36

文章被收录于专栏：Python七号

你好，我是 somenzz，可以叫我征哥，今天分享两种自定义上下文管理器的方法，并比较它们的性能。

上下文管理器相信你也用过，就是 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 次，均是方法二更快：