你好,我是 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) 的优雅和便捷。
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。