Python - 上下文管理

如果需要精确的分配和释放资源,就需要上下文管理 context manager ,上下文管理器在 Python 中最常见的例子就是 with 语句

有状态资源的管理

很多时候我们需要操作一些有状态的资源,例如 类文件、套接字等

通常情况下我们会 建立连接 —> 做一些事 —> 释放连接

f = open('aaaa', 'w')
f.write('hello, world')
f.close()

import socket

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect(('hello.org', 20333))
s.sendall('hello, world')
print(s.recv(1024))

如果没有 close ,只有程序中断或者服务端主动关闭才会释放资源,否则会一直占用资源,有兴趣的朋友可以试试下面的代码

files = []
file = open('aaaa', 'w')
file.write('aaaaa')
file.close()
for i in range(1000009):
files.append(open('aaaa', 'w'))

在 macOS/Linux 上会抛出 OS Error 的异常,因为打开的 file descriptor 超过限制了,Windows 估计会蓝屏 😀

OSError: [Errno 24] Too many open files

try..finally 清理资源

我们也可以通过 try..finally 语句来实现对资源的清理,无论是否触发异常,最后都将清理资源

try:
f = open('w', 'w')
f.write('hello, world')
except Exception as e:
raise e
finally:
f.close()

with 语句

如果看过一些 Python 的教程或者书籍的朋友可能会了解,操作文件和套接字都建议使用 with 语句

with open('aaa', 'w') as f:
f.write('hello, world')

>>> f.closed
True

with 语句下操作文件的代码非常的简洁,并且能够自动关闭文件,这是因为 TextIOWrapper 这个类实现了 __enter____exit__ 这两个魔术方法,也就是所谓的上下文管理

我们也可以自己实现一个 Open 类,使其实现上下文管理

class Open:
def __init__(self, file, mode):
self.open_file = open(file, mode)

def __enter__(self):
return self.open_file

def __exit__(self, *exception):
self.open_file.close()

>>> with Open('aaaa', 'w') as f:
f.write('aaaa')

>>> f.closed
True

__enter__ 方法是在 with 初始化时调用,__exit__ 方法在所有语句结束后调用,可以用来处理异常和关闭资源

异常处理

上下文管理器根据__exit__ 方法的返回值来决定是否抛出异常,如果没有返回值或者返回值为 False ,则异常由上下文管理器处理,如果为 True 则由用户自己处理

__exit__ 接受三个参数,exception_type、exception_value、traceback,我们可以根据这些值来决定是否处理异常

下面这个例子,捕捉了 AttributeError 的异常,并打印出警告

class Open:
def __init__(self, file, mode):
self.open_file = open(file, mode)

def __enter__(self):
return self.open_file

def __exit__(self, type, value, tb):
self.open_file.close()
if type is AttributeError:
print('handing some exception')
return True

试着故意打错方法的名称,并没有引发异常

>>> with Open('aaa', 'w') as f:
f.writeee('aaaa')
handing some exception

使用 contextmanager

由于上下文管理非常有用,Python 中有一个专门用于实现上下文管理的标准库,这就是 contextlib

有了 contextlib 创建上下文管理的最好方式就是使用 contextmanager 装饰器,通过 contextmanager 装饰一个生成器函数,yield 语句前面的部分被认为是 __enter__ 方法的代码,后面的部分被认为是 __exit__ 方法的代码

from contextlib import contextmanager

@contextmanager
def file(path, mode):
open_file = open(path, mode)
yield open_file
open_file.close()

通过上面这种方法,我们就能够简单的创建一个支持上下文管理的函数

contextmanager 的简单实现

我之前实现过一个非常粗糙的 contextmanager 装饰器

class ContextManager:
def __init__(self, fn):
self.fn = fn

def __enter__(self):
return next(self.fn)

def __exit__(self, exc_type, exc_val, exc_tb):
try:
next(self.fn)
except StopIteration:
pass


def cm(fn):
def wrapper(*args, **kwargs):
return ContextManager(fn(*args, **kwargs))

return wrapper


@cm
def open_file(path, mode):
open_file = open(path, mode)
yield open_file
open_file.close()