with语句详解 2 —— contextlib模块提供上下文管理工具

相关链接：
1. with 语句的常见用法
2. contextlib模块提供上下文管理工具
3. 异步上下文管理器

抽象类：class contextlib.AbstractContextManager

一个为实现了object.__enter__()与object.__exit__()的类提供的abstract base class。为object.__enter__()提供的一个默认实现是返回self，而object.__exit__()是一个默认返回None的抽象方法。

一、核心装饰器工具：@contextmanager 装饰器

通过生成器函数快速创建上下文管理器，无需手动实现 __enter__ 和 __exit__ 方法。生成器中 yield 前为 __enter__ 逻辑，yield 后为 __exit__ 逻辑。

示例1：资源状态管理

from contextlib import contextmanager

@contextmanager
def custom_context():
    print("进入上下文")
    yield "资源对象"
    print("退出上下文")
with custom_context() as res:
    print(f"操作资源： {res}")

输出：

进入上下文
操作资源： 资源对象
退出上下文

示例2：文件操作管理

from contextlib import contextmanager

@contextmanager
def file_opener(file_path, mode):
    file = open(file_path, mode)  # 初始化资源（类似 __enter__）
    try:
        yield file                # 返回资源对象供 with 块使用
    finally:
        file.close()              # 清理资源（类似 __exit__）‌

# 使用示例
with file_opener("test.txt", "w") as f:
    f.write("Hello World")

二、资源管理工具：closing() 函数

自动调用对象的 close() 方法，适用于需显式关闭的资源（如网络连接、文件对象）。
示例：

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen("https://example.com")) as page:
    content = page.read()

三、异常处理工具：suppress() 函数

忽略指定异常，避免冗余的 try-except 结构。
示例：

from contextlib import suppress
import os

print("忽略指定异常 1")
with suppress(FileNotFoundError):
    print('os.remove("不存在的文件.txt")')
    os.remove("不存在的文件.txt")
print("忽略指定异常 2")
os.remove("不存在的文件.txt")

输出：

忽略指定异常 1
os.remove("不存在的文件.txt")
忽略指定异常 2
Traceback (most recent call last):
  File "<input>", line 10, in <module>
FileNotFoundError: [WinError 2] 系统找不到指定的文件。: '不存在的文件.txt'

四、输出重定向工具：redirect_stdout() 和 redirect_stderr()

临时将标准输出或错误重定向到指定对象（如文件或内存缓冲区）。

示例：redirect_stdout()

from contextlib import redirect_stdout
import io

with redirect_stdout(io.StringIO()) as f:
    print("此输出被重定向")
    help(pow)
print(f"捕获内容： {f.getvalue()}")

输出：

捕获内容： 此输出被重定向
Help on built-in function pow in module builtins:
pow(base, exp, mod=None)
    Equivalent to base**exp with 2 arguments or base**exp % mod with 3 arguments
    Some types, such as ints, are able to use a more efficient algorithm when
    invoked using the three argument form.

示例：redirect_stderr()

from contextlib import redirect_stderr
import io
import sys

with redirect_stderr(io.StringIO()) as f:
    print("此输出被重定向",file=sys.stderr)
    try:
        os.remove("不存在的文件.txt")
    except FileNotFoundError as e:
        print(e,file=sys.stderr)

print(f"捕获内容： {f.getvalue()}")

输出：

捕获内容： 此输出被重定向
[WinError 2] 系统找不到指定的文件。: '不存在的文件.txt'

五、临时改变工作目录工具：contextlib.chdir

它用于临时改变当前工作目录。使用这个函数可以避免在代码执行完毕后需要手动恢复原始工作目录的麻烦，因为它会在代码块执行完毕后自动恢复到原来的工作目录。
示例：

import contextlib
import os

# 假设我们想临时改变当前工作目录
new_path = '../links'
print("进入前工作目录：", os.getcwd())
# 使用 contextlib.chdir
with contextlib.chdir(new_path):
    # 在这个 with 块中，当前工作目录被临时改变为 new_path
    print("当前工作目录：", os.getcwd())
    # 在这里执行一些依赖于当前工作目录的操作

# with 块执行完毕后，自动恢复到原来的工作目录
print("恢复后的工作目录：", os.getcwd())

输出：

进入前工作目录： D:\FluentPython2\Test
当前工作目录： D:\FluentPython2\links
恢复后的工作目录： D:\FluentPython2\Test

六、使上下文管理器能用作装饰器的基类：contextlib.ContextDecorator

与往常一样，继承自 ContextDecorator 的上下文管理器必须实现 __enter__ 与 __exit__ 。即使用作装饰器， __exit__ 依旧会保持可能的异常处理。
示例：

from contextlib import ContextDecorator

class mycontext(ContextDecorator):
    def __enter__(self):
        print('Starting')
        return self

    def __exit__(self, *exc):
        print('Finishing')
        return False

@mycontext()
def function():
    print('function()函数')
 
print(f'function(){5*"="}')
function()
print(f'with {5*"="}')
with mycontext():
    print("with代码块")

输出：

function()=====
Starting
function()函数
Finishing
with =====
Starting
with代码块
Finishing

七、使用class contextlib.ExitStack管理多个上下文

该上下文管理器的设计目标是使得在编码中组合其他上下文管理器和清理函数更加容易，尤其是那些可选的或由输入数据驱动的上下文管理器。

1. enter_context(cm)

进入一个新的上下文管理器并将其 __exit__() 方法添加到回调栈中。 返回值是该上下文管理器自己的 __enter__() 方法的输出结果。

示例，通过一个如下的 with 语句可以很容易处理一组文件：

with contextlib.ExitStack() as stack:
   files = [stack.enter_context(open(fname)) for fname in filenames]
   # 所有已打开的文件都将在 with 语句结束时
   # 自动被关闭，即使此后打开列表中文件的
   # 尝试引发了异常

2. push(exit)

将一个上下文管理器的 __exit__() 方法添加到回调栈。
示例：

class ClsContext:
    def __enter__(self):
        print("上下文管理器 __enter__()")
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"上下文管理器 __exit__() exc_type={exc_type};exc_val={exc_val};exc_tb={exc_tb}")
        return True

with contextlib.ExitStack() as stack:
    clsContext = stack.push(ClsContext())
    print("上下文管理器 ExitStack() 代码块")

输出：

上下文管理器 ExitStack() 代码块
上下文管理器 __exit__() exc_type=None;exc_val=None;exc_tb=None

如果传入了一个不是上下文管理器的对象，此方法将假定它是一个具有与上下文管理器的 __exit__() 方法相同的签名的回调并会直接将其添加到回调栈中。
示例：

def contextFunction(exc_type, exc_val, exc_tb):
    print(f"上下文管理器contextFunction __exit__() exc_type={exc_type};exc_val={exc_val};exc_tb={exc_tb}")

with contextlib.ExitStack() as stack:
    stack.push(contextFunction)
    print("上下文管理器 ExitStack() 代码块")

输出：

上下文管理器 ExitStack() 代码块
上下文管理器contextFunction __exit__() exc_type=None;exc_val=None;exc_tb=None

3. callback(callback, /, *args, **kwds)

• 接受一个任意的回调函数和参数并将其添加到回调栈。
• 与其他方法不同，通过此方式添加的回调无法屏蔽异常（因为异常的细节不会被传递给它们）。
• 传入的回调将被该函数返回，允许此方法作为函数装饰器使用。
  示例：

import contextlib

with contextlib.ExitStack() as stack:
    # 注册一个回调函数来执行清理工作
    stack.callback(lambda: print("Cleaning up..."))
    print("上下文管理器 ExitStack() 代码块")

输出：

上下文管理器 ExitStack() 代码块
Cleaning up...

4. pop_all()

将回调栈传输到一个新的 ExitStack 实例并返回它。 此操作不会唤起任何回调 —— 作为替代，现在当新栈被关闭时它们将（显式地或是在一条 with 语句结束时隐式地）被唤起。

示例，一组文件可以像下面这样以“一个都不能少”的操作方式被打开:

with contextlib.ExitStack() as stack:
    files = [stack.enter_context(open(fname)) for fname in filenames]
    # 持有 close 方法，但暂时不调用它。
    close_files = stack.pop_all().close
    # 如果任何文件打开失败，则所有之前打开的文件
    # 将自动被关闭。 如果所有文件都被成功地打开，
    # 即使在 with 语句结束后它们仍将保持打开状态。
    # 随后可以显式唤起 close_files() 来全部关闭它们。

5. close()

立即展开回调栈，按注册时的相反顺序唤起其中的回调。 对于任何已注册的上下文管理器和退出回调，传入的参数将表明没有发生异常。

<< 上一篇： 1. with 语句的常见用法
下一篇 >>： 3. 异步上下文管理器