Python contextlib 模块

Python 标准库 contextlib 提供了工具便于在 with 语句中使用

1. 忽略异常

简单示例

import contextlib

with contextlib.suppress(ZeroDivisionError):
    print("before x = 1 / 0")
    x = 1 / 0
    print("after x = 1 / 0")

print("exit x = 1 / 0")

打印

before x = 1 / 0
exit x = 1 / 0

其等价于:

try:
    print("before x = 1 / 0")
    x = 1 / 0
    print("after x = 1 / 0")
except ZeroDivisionError:
    pass

print("exit x = 1 / 0")

使用 contextlib.suppress 的写法更简单优雅

2. 临时重定向

import contextlib
import io

temp = io.StringIO("Python")
with contextlib.redirect_stdout(temp):
    print("Test", end="!")

print(temp.getvalue())

结果

Test!n

io.StringIO 相当于一个保存在内存的文件流对象, 初始内容为 “Python”, 文件指针仍指向文件开头, 其后写入内容时呈现覆盖效果

第一个 print 将文本写入重定向位置 temp, 第二个 print 打印到终端(重定向被取消)

除了重定向标准输出, 还可重定向标准错误 contextlib.redirect_stderr

3. 创建环境管理器

一般创建环境管理器需要借助于类且实现 __enter__ 和 __exit__ 方法, contextlib.contextmanager 以装饰器的形式提供另一种创建环境管理器的方法

import contextlib

@contextlib.contextmanager
def fun_1(s):
    print("<h1>")
    yield s
    print("</h1>")


with fun_1("Python") as content:
    print(content)

打印

<h1>
Python
</h1>

fun_1 生成器函数仅能消耗一次, 也就是只能 yield 一次

执行流程:

1. 执行 fun_1("Python")(已被装饰器修饰) 返回一个类
2. 执行该类的 __enter__ 方法, 内部调用 next(fun_1)
3. 执行 print(content)
4. 执行该类的 __exit__ 方法, 内部仍调用 next(fun_1), 如果生成器函数 fun_1 不退出, 抛出 RuntimeError 异常

错误示例

@contextlib.contextmanager
def fun_2(s):
    print("<h1>")
    yield s
    yield s
    print("</h1>")


with fun_2("Python") as content:
    print(content)

抛出异常 RuntimeError: generator didn't stop

MySQL 连接应用

import MySQLdb

@contextlib.contextmanager
def client():
    conn = MySQLdb.connect()
    try:
        yield conn
    finally:
        conn.close()


with client() as conn:
    pass

如果对象的释放函数为 close, 则可以使用标准库提供的 contextlib.closing

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen('http://www.python.org')) as page:
    for line in page:
        print(line)