0
点赞
收藏
分享

微信扫一扫

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



你好,我是 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 次,均是方法二更快:

两种方法自定义Python上下文管理器-contextlib_反射

最后的话

进入代码前需要做一些预处理,离开后需要再次处理的时候,上下文管理器就可以出场了,如果这样的处理变多的时候,你会发现 contextlib 真的是个好东西,大大减轻了程序员的心智和记忆负担。本文分享了两种实现 contextlib 的方法,推荐使用类来实现 contextlib,希望对你有所帮助。



举报

相关推荐

0 条评论