Python中with语句详解, 2种方式创建自定义的上下文管理器

1. Python中的资源管理

在编程中面临的一个常见问题是如何正确管理外部资源，例如文件、锁和网络连接。如果创建和打开资源后, 而不实现关闭, 就会出现内存泄漏.

正确管理资源需要一个Setup阶段和一个Teardown, Teardown阶段需要执行一些清理操作，例如关闭文件、释放锁或关闭网络连接.

file = open("hello.txt", "w")
file.write("Hello, World!")
file.close()

如果调用 .write() 期间发生异常，.close() 代码就不会被调用, 无法保证文件被正确关闭.

在 Python 中，您可以使用两种通用方法来处理资源管理。您可以将代码包装在：

1. 一个 try ... finally 代码块
2. 一个 with 代码块

1.1 😐 使用 try ... finally 

使用 try ... finally 语句可以确保, .write()期间发生异常,  close() 代码也会被执行.  

file = open("hello.txt", "w")

try:
    file.write("Hello, World!")
finally:
    file.close()

1.2 🙂 使用 with 语句

Python中的 with 语句会创建一个上下文, 与 try ... finally 结构相比，with 语句可以使的代码更清晰、更安全和可重用。

使用 with 语句来实现打开 hello.txt 文件的代码如下:

with open("hello.txt", mode="w") as file:
    file.write("Hello, World!")

with 语句的通用语法如下: 

expression 必须返回一个实现上下文管理协议的对象

上下文管理协议需要实现两个特殊方法, 分别对应Setup和Teardown: 

1. .enter()被 with 语句调用以进入运行时上下文。
2. .exit()当执行离开 with 代码块时调用。

1.2.1 多个上下文管理器

Python 3支持多个上下文管理器, 在 with 语句后面可以使用逗号分隔多个上下文管理器:

with A() as a, B() as b, C() as c:
    do_somethind()

 比如操作多个文件时, 尤为有用. 比如如下代码, 将input.txt文件中的数据, 写入到output.txt中.

with open("input.txt") as input_file, open("output.txt", "w") as output_file:
    data = input_file.read()
    output_file.write(data)

2. 自定义上下文管理器

2.1 创建基于类的上下文管理器

可以通过创建一个实现了 __enter__() 和 __exit__() 的特殊方法的类, 来创建基于类的上下文管理器.

• .__enter__(self)   此方法在进入 with 上下文时调用。它的返回值绑定到 as 后的目标变量。
• .__exit__(self, exc_type, exc_value, exc_tb)    此方法在执行流离开 with 上下文时调用。如果发生异常，则exc_type、exc_value和 分别exc_tb保存异常类型、值和回溯信息。

In [1]: class MyContext:
   ...:     def __enter__(self):
   ...:         print('Entering my context')
   ...:         return [1, 2, 3]
   ...:     def __exit__(self, exc_type, exc_value, exc_tb):
   ...:         print('Leaving my context')
   ...:         print(exc_type, exc_value, exc_tb, sep='_____')

In [2]: with MyContext() as data:
   ...:     print(data[0])
   ...:
Entering my context
1
Leaving my context
None_____None_____None

In [3]: with MyContext() as data:
   ...:     print(data[3])
Entering my context
Leaving my context
<class 'IndexError'>_____list index out of range_____<traceback object at 0x0000021814964A08>
---------------------------------------------------------------------------
IndexError: list index out of range

可以看到, 使用 with 语句之后,  执行顺序为 

① data = MyContext().__enter__()  ,  打印"Entering my context" 到输出.

② print(data[0]),  打印 1 到输出

③ MyContext().__exit__() , 打印"Leaving my context, None_____None_____None"到输出

如果 with 内的代码块内发生异常,  依然会执行__exit__() 中的代码.

2.2 创建基于函数的上下文管理器

还可以使用contextlib标准库中的contextmanager函数来创建自定义的基于函数的上下文管理.

使用@contextmanager生成器函数创建上下文管理器的代码样例：

In [1]: from contextlib import contextmanager

In [2]: @contextmanager
   ...: def my_context():
   ...:     print('Entering my context')
   ...:     yield [1, 2, 3]
   ...:     print('Leaving my context')

In [4]: with my_context() as data:
   ...:     print(data[0])
Entering my context
1
Leaving my context

In [5]: with my_context() as data:
   ...:     print(data[3])
Entering my context
---------------------------------------------------------------------------
IndexError: list index out of range

与一般函数不同, @contextmanager装饰过的函数使用 yield 返回数据, 结果赋值给 with 语句中 as 的目标变量. yield 语句之前的代码为 Setup 阶段的执行代码.   yield 语句之后的代码为 Teardown 阶段的代码.

与使用类实现的上下文管理器不同的事,  contextmanager不自动进行异常捕获, 若 with 中的代码出现异常, 无法保证 Teardown阶段的代码的执行. 如果需要处理异常, 需要写 try ... catch 语句.

In [6]: @contextmanager
   ...: def my_context():
   ...:     print('Entering my context')
   ...:     try:
   ...:         yield [1, 2, 3]
   ...:     except IndexError as e:
   ...:         print('Err Msg: ', e)
   ...:     finally:
   ...:         print('Leaving my context')

In [7]: with my_context() as data:
   ...:     print(data[3])
Entering my context
Err Msg:  list index out of range
Leaving my context

总结:

① 什么是Python with 语句 以及如何使用

② 两种实现自定义上下文的方法