Python3 with语句和上下文管理器

0

with语句和上下文管理器

with语句

在 Python 中,传统的处理文本文件读写的方式都是这样:

file = open('1.txt', ...)
...  # 读写操作
file.close()

这也是多数编程语言都使用的常规的处理文件的流程:打开文件、读写文件、关闭文件。

不过这种文本文件的处理方式存在一些潜在的问题:对于复杂的读写操作,可能在其中的某一步出现异常,造成 file.close() 语句未能顺利执行;如果读写操作过于复杂,也可能忘记添加 file.close() 语句。这几种情况下,文件都可能因为未被关闭而保存失败,或资源持续占用等问题。

文件是否关闭可以通过编写测试代码来检查,智能的 IDE 也可能发现未关闭的文件。对于可能出现错误的情况,一般采取的解决方法是将其放在 try 语句中捕获可能的错误,并在 finally 中确保关闭文件:

file = open('1.txt', ...)
try:
    ...  # 读写操作
except:
    ...  # 错误处理
finally:
    file.close()

以上做法稍显复杂。实际上 Python 的语法中存在一种 with 语句,也称为上下文管理语句。使用上下文管理语句操作文件时,其代码为:

with open('1.txt', ...) as file:
    ...  # 读写操作

这时,不需要手动在代码中关闭文件,因为 with 语句相当于打开了一个上下文,只要当 with 内语句执行完毕或出现异常时,它都会自动执行一条退出命令,即 file.close()

with 语句的结构为:

with context_manager as cm:
    ...  # some operations

以上语句通过 with 语句生成了一个上下文(context),并可能通过 as 语句得到上下文管理对象。在 with 语句中的所有操作,都位于这个上下文以内。不管这些操作成功还是失败,它都需要进入和退出这个环境。而它退出时,便会进行相应的清理工作:对于文件操作,对应的会执行 file.close() 语句。

以上便是 Python 内置的关于文件的上下文管理操作。上下文管理操作是 Python 的一个高级语法,在合理的场景使用,不仅降低了程序出错的概率,而且可以使程序更简洁明了。

另一个使用上下文管理器的典型场景是 Python 中的多线程锁。使用 with 语句配合多线程锁时,锁的获取与释放都会自动执行:

import threading

lock = threading.Lock()
with lock:
    ...

接下来看看如何实现自定义的上下文管理操作。

上下文管理器

可以自行实现一个上下文管理器(context manager)来完成上下文管理操作。

如果一个对象是一个上下文管理器,它需要实现两个特殊方法:.__enter__().__exit__()

.__enter__() 方法是进入上下文语句后需要进行的操作,例如它可能会初始化一些上下文管理所需的对象。该方法可以返回一个合适的上下文管理对象,该对象会被 as 部分接收,并在上下文中被处理。

.__exit__() 方法是退出上下文时需要进行的操作,该方法的完整结构为:.__exit__(self, exc_type, exc_val, exc_tb) ,它需要四个参数,后三个参数分别是异常类型、异常信息和回溯栈。在该方法中,可以释放所需资源,或处理可能遇到的异常等。

以下示例是一个自定义的上下文管理器:

import sys

class Log:
    def __init__(self, filename):
        self.__file = open(filename, 'a', encoding='utf-8')
    def __enter__(self):
        self.__saved_stdout = sys.stdout
        sys.stdout = self.__file
        return self.__file
    def __exit__(self, exc_type, exc_value, exc_tb):
        self.__file.close()
        sys.stdout = self.__saved_stdout

该示例可以比较好的说明上下文管理器的一些性质。一个调用该上下文管理器的示例为:

from time import ctime

print('hello')
with Log('log.txt') as log:
    print(f'[{ctime()}] sqrt of 10 is: {10 ** 0.5}')
print('world')

运行该程序,在观察到输出的同时,还会在本地生成一个日志文件:

$ python -u with.py hello world $ cat log.txt [Wed Jul 13 10:59:34 2022] sqrt of 10 is: 3.1622776601683795

首先在 .__init__() 方法中,通过初始化该类时传递的参数,打开了一个文件对象;在 .__enter__() 方法中,做了一些初始化工作:将系统的标准输出流重定向到初始化打开的文件对象,并返回该文件对象用于上下文操作;在 .__exit__() 方法中,保存该日志文件,并改回系统的标准输出流。

如果在执行 with 语句的语句体期间发生了异常,则会立即进入 .__exit__() 方法中,确保文件正常关闭且标准输出流被改回,同时其余的 3 个参数会给出异常提示。在其它情况下,这三个参数均为 None

contextlib标准库:上下文管理工具

contextlib 是 Python3 标准库之一,它提供了一些上下文处理的工具,帮助编写合适的上下文管理器来处理 with 语句。

装饰器形式的上下文管理器

ContextDecorator 基类用于将一个上下文管理器变成一个装饰器,通过继承该类,可以让一个上下文管理器装饰函数,使函数变成上下文。例如:

from contextlib import ContextDecorator

class Log(ContextDecorator):
    ...

@Log('log.txt')
def func():
    ...  # some operations
    print('... record ...')

该函数内的代码便会被作为上下文被执行。这种装饰器形式可以不显式使用 with ,不过它也没有 as ,不能利用得到的上下文管理对象。


如果每次定义一个上下文管理器都需要创建一个新类,不免会有些麻烦。contextlib 提供了一种函数形式的更简单的自定义上下文管理器的方式:通过使用 @contextmanager 装饰器,可以将一个函数变成上下文管理器。被装饰的函数需要有类似如下的结构:

@contextmanager
def func_name(...):
    cm = ...  # get context object
    try:
        yield cm
    finally:
        ...  # exit context

在这个函数(或者说生成器)形式的上下文管理器中,yield 生成的内容将会作为 as 的目标值,而 finally 对应的语句将会是退出上下文语句后执行的部分。

以下是一个上下文管理函数示例:

import pymysql
from contextlib import contextmanager

@contextmanager
def open_database(db_name):
    connection = pymysql.Connect(
        host=settings.db_host,
        user=settings.db_username,
        password=settings.db_password,
        database=db_name
    )
    try:
        yield connection.cursor()
    finally:
        connection.commit()
        connection.close()

在这个上下文管理函数内,在进入上下文时连接到数据库中,并返回一个游标作为上下文对象;借助游标可以对数据库做一些做一些操作,例如:

with open_database('demo') as db:
    db.execute(
        'INSERT INTO user(username, password) VALUES (%s, %s)',
        [input(), input()]
    )

在退出上下文时,将操作更新到数据库中,并断开连接。这种函数形式的上下文管理器显然更简洁,不用特地编写一个类。

其它上下文管理器

closing(thing) 是一个简单的上下文管理器,它可以将任意对象变成一个上下文管理对象,并在结束上下文时调用其 .close() 方法。

许多对象都包含 .close() 方法,就可以通过该函数将其变成上下文管理对象,例如套接字:

import socket
from contextlib import closing

with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as client:
    client.connect(('192.168.1.105', 8000))
    client.send(message.encode())

supress(*exceptions) 用于在一个上下文中忽略特定类型的错误,相当于将其放入一个 try 内并捕获特定错误。

例如,可以使用该上下文管理器在创建文件或目录时忽略文件已存在的问题,从而无需提判断文件是否存在:

from contextlib import suppress
from os import mkdir

with suppress(FileExistsError):
    mkdir('imgs')

还有两个上下文管理器 redirect_stdout(io)redirect_stderr(io) 可以重定向标准输出流和标准错误流,其原理和用法和以上编写的 Log 类似。

除此之外,还有一类特殊的异步上下文管理器,它们相比上文介绍的上下文管理器在命名时前面都带了个“ a ”(包括需要实现的特殊方法也是 .__aenter__().__aexit__() )。由于异步上下文管理器主要应用于异步编程中,除此之外和普通的上下文管理器相差不大,因此这里不再介绍。

参考资料/延伸阅读

https://docs.python.org/3/library/contextlib.html
Python3 标准库官方文档—— contextlib

https://docs.python.org/3/reference/datamodel.html#context-managers
Python3 复合语句 with ,里面还有两个链接指向相关话题

https://peps.python.org/pep-0343/
PEP343 ——关于 with 语句的提案和背景

京ICP备2021034974号
contact me by hello@frozencandles.fun