Python 上下文管理器

1. Abstract

在 Python 中读写文件，我们都用 with open(filename) as f:，这样不用担心忘记关闭文件流，因为离开 with 环境时，文件会自动关闭。这东西叫上下文管理器，现在让让我们学习一下这个神奇的家伙。

2. 简介

首先推荐博文《with open为什么会自动关闭文件流》，讲的比较简单易懂，我就是看这篇博文学习上下文管理器的。我先把他的代码搬过来：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
你是否想过一个问题，打开文件会抛出异常，通常打开文件后也需要关闭文件流，为什么用 with open()语句可以不用手动关闭文件流呢？
这就是上下文管理器
"""
class Sample:
	def __init__(self):
		# 首先执行这个方法
		print("__init__")

	def __enter__(self):
		# 然后会自动调用这个方法，可以理解为获取资源
		print("__enter__")
		return self

	def __exit__(self, exc_type, exc_val, exc_tb):
		# 这个函数会自动调用，当跳出 with 语句的时候，目的是为了释放资源
		print("__exit__")

	def toDo(self):
		print("to do something")

"""
__enter__和__exit__构成了上下文管理器
"""
# 这个用法是不是很像 with open()呢？
with Sample() as sample:
	print("aaa")

运行代码，输出：

3. as 是把 Sample() 对象赋值给 sample 了吗？

注意到，__enter__(self) 函数返回了一个东西：self，即 Sample() 实例的引用，进而 as sample 将 self 赋值给 sample。将 print(aaa) 改成 print(sample)，进行验证：

with Sample() as sample:
	print(sample)

##### output #####
__init__
__enter__
<__main__.Sample object at 0x000002B5FC858FD0>
__exit__

确实：sample 是一个 Sample 对象。

那应该能返回其他东西吧？再把 __enter__(self) 返回值改一改：

def __enter__(self):
	# 然后会自动调用这个方法，可以理解为获取资源
	print("__enter__")
	return 'a'  # 返回 a

##### output #####
__init__
__enter__
a
__exit__

果然输出了 a。

结论：as 是把函数 def __enter__(self): 的返回值赋给了 sample。

4. 模拟文件的上下文管理器

那我可魔改这东西了，把它变成和 open 一样的文件管理器。在这之前，我们先看一看如何验证文件是否关闭：

f = open('./nihao.txt', 'r', encoding='UTF-8')
text = f.readline()
print(text)
text = f.readline()
print(text)
f.close()

##### output #####
Hello
World

####################### 关闭文件后再读 ##########################

f = open('./nihao.txt', 'r', encoding='UTF-8')
text = f.readline()
print(text)
f.close()  # 关闭文件

text = f.readline()  # 再读
print(text)

##### output #####
ValueError: I/O operation on closed file.
Hello

文件关闭后，再读就会报错。

魔改结果（将文件操作的几个步骤放到对应的魔法函数内）：

class File(object):
	""" 管理文件的打开与关闭 """

	def __init__(self, filepath, mode='r', encoding='UTF-8'):
		self.__file = open(file=filepath, mode=mode, encoding=encoding)
		print('打开文件')

	def __enter__(self):
		print('返回文件')
		return self.__file

	def __exit__(self, exc_type, exc_val, exc_tb):
		self.__file.close()
		print('关闭文件')

进行试验：

with File('./nihao.txt', 'r', encoding='UTF-8') as f:
	text = f.readline()
	print(text)
	text = f.readline()
	print(text)

##### output #####
打开文件
返回文件
Hello
World
关闭文件

一切正常。如果在 with 环境外读文件呢？

with File('./nihao.txt', 'r', encoding='UTF-8') as f:
	text = f.readline()
	print(text)
text = f.readline()  # 移到 with 外
print(text)

##### output #####
ValueError: I/O operation on closed file.
打开文件
返回文件
Hello
关闭文件

说明离开 with 后，文件就关闭了。

结论：这段魔改代码实现了和 with open() as f 一样的功能。

5. 生成器函数式上下文管理器

把博主的代码搬过来：

# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
如何把上下文管理器更加简化一下呢？
"""
import contextlib

@contextlib.contextmanager  # 这个装饰器把下面的函数包装成上下文管理器，主要利用了 yield 的特性
def myFun(arg1):
	print("begin", arg1)  # 相当于 __enter__ 里面的代码
	yield {}  # 这里必须有个生成器
	print("finished")  # 相当于 __exit__ 里面的代码

with myFun("AAA") as my:
	print("BBB")

##### output #####
begin AAA
BBB
finished

那让我们看一看，as my 干了什么：

with myFun("AAA") as my:
	print("BBB")
	print(my)  # 打印 my

##### output #####
begin AAA
BBB
{}  # 输出了空字典，和 yield {} 对应
finished

结论：my 被赋予了生成器 yield 的东西。

5.1 再次模拟文件的上下文管理器

再用这种生成器式管理器模拟一下 with open(file) as f：

@contextlib.contextmanager  # 这个装饰器把下面的函数包装成上下文管理器，主要利用了yiele的特性
def my_file(file_path, mode='r', encoding='UTF-8'):

	print('打开文件', file_path)  # 相当于 __enter__ 里面的代码
	f = open(file=file_path, mode=mode, encoding=encoding)

	print('返回文件')
	yield f  # 这里必须有个生成器

	f.close()
	print('关闭文件')  # 相当于 __exit__ 里面的代码

与 File 实行相同的验证：

with my_file('./nihao.txt', 'r', encoding='UTF-8') as f:
	text = f.readline()
	print(text)
	text = f.readline()
	print(text)

##### output #####
打开文件 ./nihao.txt
返回文件
Hello
World
关闭文件

按预期执行了：打开、返回、读取、关闭。

with my_file('./nihao.txt', 'r', encoding='UTF-8') as f:
	text = f.readline()
	print(text)
text = f.readline()
print(text)

##### output #####
ValueError: I/O operation on closed file.
打开文件 ./nihao.txt
返回文件
Hello
关闭文件

说明文件离开 with 后，文件关闭了。

5.2 这个生成器如果是循环 yield 呢？

先看一看添加 @contextlib.contextmanager 前后，生成器有什么变化：

前

def generate(n):
	for i in range(n):
		print('before')
		yield i
		print('after')

g = generate(2)
print(next(g))

print('我在中间')

for ii in g:
	print(ii)

输出：

before
0
我在中间  # after 在后，确实暂停了函数的执行
after
before
1
after

我在中间 紧随 0 之后，说明，yield 后，确实暂停了函数的执行，而后再需要 yield 时，会执行 yield 语句后的内容：print('after')。

后

@contextlib.contextmanager
def generate(n):
	for i in range(n):
		print('before')
		yield i
		print('after')

g = generate(2)
print(next(g))

输出：

TypeError: '_GeneratorContextManager' object is not an iterator

添加 @contextlib.contextmanager 后，def generate(n): 已不再是一个 iterator，而是一个 '_GeneratorContextManager' object。

那么，对它实施 with ... as ... 呢？

with generate(2) as g:
	print(type(g))
	print(g)

输出：

RuntimeError: generator didn't stop
before
<class 'int'>  # g=0 是 int
0
after
before  # 到下一轮循环了

前面是正常的：before $\to$ 返回 $\to$ with 环境内的内容 $\to$ after。异常的是：RuntimeError: generator didn't stop，继续下一轮的 before。

结论：添加 @contextlib.contextmanager 后，生成器函数变成了 '_GeneratorContextManager' object，但只能 yield 一个东西，赋给 as ...。