上下文管理器和with语句

嗨，大家好，我是暴走的海鸽。今天，我将向大家介绍上下文管理器和with语句的使用方法。

以下上下文管理器、with未作特殊说明时均为同步原语。

什么是上下文管理器和with

有人认为Python中的with语句是一个晦涩难懂的特性，但是只要你了解了其背后的原理，就不会感觉到神秘了。with语句实际上是非常有用的特性，有助于编写清晰易读的Python代码。

而上下文管理器存在的目的便是管理with语句，就像迭代器的存在是为了管理for语句一样。

那么，究竟 with 语句要怎么用，与之相关的上下文管理器（context manager）是什么，它们之间又有着怎样的联系呢？

在任何一门编程语言中，文件的输入输出、数据库的连接断开等，都是很常见的资源管理操作。但资源都是有限的，在写程序时，我们必须保证这些资源在使用过后得到释放，不然就容易造成资源泄露，轻者使得系统处理缓慢，重则会使系统崩溃。

在 Python 中，上下文管理器（Context Manager）是一种资源管理的机制，用于在使用资源（如文件、网络连接、数据库连接等）之前分配资源，然后在使用完毕后释放资源。with 语句是用来简化上下文管理器的使用的语法糖，使代码更加清晰和简洁。

一个典型的案例便是内置的open()函数

使用try/finally结构保证文件的最终关闭：

f = open('test.txt', 'w')
try:
    f.write('hello')
finally:
    f.close()

然而，这一套固定模板较为冗长，还容易漏写，代码可读性也较差，因此我们一般更倾向于with语句。

with open('test.txt', 'w') as f: 
    f.write('hello')

另外一个典型的例子，是 Python 中的 threading.lock 类。

比如我想要获取一个锁，执行相应的操作，完成后再释放，那么代码就可以写成下面这样：

some_lock = threading.Lock()
some_lock.acquire()
try:
    ...
finally:
    some_lock.release()

而对应的 with 语句，同样非常简洁：

some_lock = threading.Lock()
with somelock:
    ...

什么场景建议考虑使用上下文管理器和with语句

从前面的例子中可以看出，上下文管理器的常见用途是自动打开和关闭文件以及锁的申请和释放。

不过，你还可以在许多其他情况下使用上下文管理器，综合来看有以下几种场景：

1) 打开-关闭

• 如果想自动打开和关闭资源，可以使用上下文管理器。
• 例如，可以使用上下文管理器打开一个套接字并关闭它。

2) 锁定-释放

• 上下文管理器可以帮助你更有效地管理对象的锁。它们允许你获取锁并自动释放锁。

3) 启动-停止

• 上下文管理器还能帮助您处理需要启动和停止阶段的场景。
• 例如，您可以使用上下文管理器启动计时器并自动停止。

4) 更改 - 重置

• 上下文管理器可以处理更改和重置场景。

• 例如，您的应用程序需要连接多个数据源。它有一个默认连接。
  要连接到另一个数据源：

  • 首先，使用上下文管理器将默认连接更改为新连接。
  • 第二，使用新连接。
  • 第三，完成新连接的操作后，将其重置回默认连接。

5) 进入 - 退出

• 上下文管理器可以处理进入和退出的场景。

常见的上下文管理器和with语句使用

为了更充分地掌握和运用这一特性，一个较为有效的途径是深入研究该特性在各个广泛应用的场景中的实际运用。

上下文管理器和 with 语句在 Python 中的使用非常广泛，涵盖了许多不同的场景。以下是一些常见的上下文管理器和 with 语句的使用场景：

with open("file.txt", "r") as file:
    content = file.read()

import sqlite3

with sqlite3.connect("mydatabase.db") as connection:
    cursor = connection.cursor()
    cursor.execute("SELECT * FROM mytable")

 import requests

 url = 'https://www.example.com'

 # 使用Session对象作为上下文管理器
 with requests.Session() as session:
     # 发起GET请求
     response = session.get(url)

     # 在此处处理响应
     print(response.status_code)
     print(response.text)

 # 在退出上下文时，底层连接会被关闭

import threading

lock = threading.Lock()

with lock:
    # 执行需要线程同步的代码

import fcntl

with open("file.txt", "r") as file:
    fcntl.flock(file, fcntl.LOCK_EX)
    # 执行需要文件锁的代码
    fcntl.flock(file, fcntl.LOCK_UN)

import time

class TimerContextManager:
    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        elapsed_time = self.end_time - self.start_time
        print(f"Code executed in {elapsed_time} seconds")

with TimerContextManager():
    # 执行需要测量时间的代码

class TestContextManager:
    def __enter__(self):
        # 在测试开始时分配资源
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # 在测试结束时释放资源

with TestContextManager() as test:
    # 执行测试代码

这些只是很小的一部分示例，上下文管理器和 with 语句的使用可以根据具体的应用场景而变化。总体而言，它们是一种确保资源正确管理的强大方式，提高了代码的可读性和可维护性。

上下文管理器和with的本质

上面我们了解了一些常见的上下文管理器和with的用法，有内置的上下文管理器和自定义的上下文管理器。可能你已经从中看出点端倪了，不过不妨让我们深入的了解下这一特性的本质。

上下文管理器的实现

一个上下文管理器必须实现两个方法：__enter__ 和 __exit__。

• __enter__(self)：在进入 with 语句块时被调用，返回资源对象，该对象会被赋值给 as 后面的变量。
• __exit__(self, exc_type, exc_value, traceback)：在离开 with 语句块时被调用，用于资源的释放。如果 with 语句块中发生异常，异常的信息（异常类型（如ZeroDivisionError）、异常示例（如错误消息）、traceback对象）将作为参数传递给 __exit__ 方法。如果__exit__方法返回None或者True之外的值，with块中任何异常都会向上冒泡。

在try/finally语句中调用sys.exec_info()，得到的便是__exit__接收的三个参数。因此可以在__exit__中对异常进行处理。

我们来看一个发生异常的例子，观察 __exit__ 方法拿到的异常信息是怎样的：

class TestContext:
    def __enter__(self):
        print("__enter__")
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        print("exc_type:", exc_type)
        print("exc_value:", exc_value)
        print("exc_tb:", exc_tb)
        print("__exit__")


with TestContext() as t:
    # 这里会发生异常
    a = 1 / 0
    print('t: %s' % t)

"""
__enter__
exc_type: 
exc_value: division by zero
exc_tb: 
__exit__
Traceback (most recent call last):
  File "/mnt/d/jinzhuan/rss/tpl_hf.py", line 15, in 
    a = 1 / 0
ZeroDivisionError: division by zero
"""

可以在__exit__处理异常，这也是很多框架中用于消除大量try/exception/finally的手段之一，比如django中就有这样处理数据库操作异常时的代码，见文末源码。

with 语句

直接调用__enter__和__exit__方法是使用上下文管理器的方式之一，但是with 语句是一种更加简洁的处理上下文管理器的方式。它的语法如下：

with expression as variable:
    # code block

• expression 返回一个上下文管理器对象。
• variable 是一个变量，用于存储 expression 返回的上下文管理器对象。

在 with 语句块中，上下文管理器的 __enter__ 方法被调用，返回的对象赋值给 variable。当退出 with 语句块时，无论是正常退出还是因为异常，上下文管理器的 __exit__ 方法都会被调用。

with TimerContextManager() as timer:
    # 在这里写你想要测量执行时间的代码
    for _ in range(1000000):
        pass

# 在这里 timer 对象不再可用，执行时间已经在 __exit__ 方法中输出

使用 with 语句可以确保资源的正确分配和释放，使代码更加简洁和可读。上下文管理器在文件操作、数据库连接、网络连接等场景中经常被使用。在 Python 标准库中，有一些内置的上下文管理器，比如 open() 函数返回的文件对象就是一个上下文管理器。

python 3.10后支持带括号的上下文管理器

with(  
    open("text1.txt", encoding="utf-8") as f1,  
    open("text2.txt", encoding="utf-8") as f2  
):  
    print(f1.read(), f2.read())

异步上下文管理器和async with的本质

异步上下文管理器

要创建异步上下文管理器，您需要定义.__aenter__()和.__aexit__()方法。

import aiohttp
import asyncio


class AsyncSession:
    def __init__(self, url):
        self._url = url

    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        response = await self.session.get(self._url)
        return response

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        await self.session.close()


async def check(url):
    async with AsyncSession(url) as response:
        print(f"{url}: status -> {response.status}")
        html = await response.text()
        print(f"{url}: type -> {html[:17].strip()}")


async def main():
    await asyncio.gather(
        check("https://realpython.com"),
        check("https://pycoders.com"),
    )


asyncio.run(main())

# https://realpython.com: status -> 200
# https://pycoders.com: status -> 200
# https://pycoders.com: type -> 
# https://realpython.com: type ->  200
# https://realpython.com: type ->  200
# https://pycoders.com: type -> 

contextlib包简化上下文管理器的实现和管理

contextlib 模块提供了一些实用工具，用于简化创建和操作上下文管理器。以下是一些常用的 contextlib 工具：

from contextlib import ExitStack

def example_function(file1_path, file2_path):
    with ExitStack() as stack:
        # 打开第一个文件
        file1 = stack.enter_context(open(file1_path, 'r'))
    
        # 打开第二个文件
        file2 = stack.enter_context(open(file2_path, 'r'))

        # 在此处进行文件处理，无需手动关闭文件

# 使用例子
example_function('file1.txt', 'file2.txt')

这些工具使得使用上下文管理器更加方便和灵活。contextlib 是 Python 标准库中一个强大而灵活的模块，用于简化上下文管理器的使用。

关于contextlib包的更多使用技巧，请阅读源码或参考官方文档。

更多应用示例

自定义文件上下文管理类

class FileManager:
    def __init__(self, name, mode):
        print('calling __init__ method')
        self.name = name
        self.mode = mode 
        self.file = None
        
    def __enter__(self):
        print('calling __enter__ method')
        self.file = open(self.name, self.mode)
        return self.file


    def __exit__(self, exc_type, exc_val, exc_tb):
        print('calling __exit__ method')
        if self.file:
            self.file.close()
            
with FileManager('test.txt', 'w') as f:
    print('ready to write to file')
    f.write('hello world')
    
## 输出
calling __init__ method
calling __enter__ method
ready to write to file
calling __exit__ method

Redis分布式锁

from contextlib import contextmanager  
  
@contextmanager  
def lock(redis, lock_key, expire):  
    try:  
        locked = redis.set(lock_key, 'locked', expire)  
        yield locked  
    finally:  
        redis.delete(lock_key)  
  
# 业务调用 with 代码块执行结束后 自动释放锁资源  
with lock(redis, 'locked', 3) as locked:  
    if not locked:  
        return  
    # do something ...

Redis事务和管道

from contextlib import contextmanager  
  
@contextmanager  
def pipeline(redis):  
    pipe = redis.pipeline()  
    try:  
        yield pipe  
        pipe.execute()  
    except Exception as exc:  
        pipe.reset()  
              
# 业务调用 with 代码块执行结束后 自动执行 execute 方法  
with pipeline(redis) as pipe:  
    pipe.set('key1', 'a', 30)  
    pipe.zadd('key2', 'a', 1)  
    pipe.sadd('key3', 'a')

让我简要解释一下这段代码：

异步文件上下文管理器

import aiofiles
import asyncio

async def read_file_async(file_path):
    async with aiofiles.open(file_path, mode='r') as file:
        content = await file.read()
        print(content)

asyncio.run(read_file_async('example.txt'))

异步HTTP请求上下文管理器

import aiohttp
import asyncio


async def fetch_data(url):
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            data = await response.text()
            print(data)

asyncio.run(fetch_data('https://baidu.com'))

django中的DatabaseErrorWrapper上下文管理器源码

class DatabaseErrorWrapper:
    """
    Context manager and decorator that reraises backend-specific database
    exceptions using Django's common wrappers.
    """

    def __init__(self, wrapper):
        """
        wrapper is a database wrapper.

        It must have a Database attribute defining PEP-249 exceptions.
        """
        self.wrapper = wrapper

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            return
        for dj_exc_type in (
                DataError,
                OperationalError,
                IntegrityError,
                InternalError,
                ProgrammingError,
                NotSupportedError,
                DatabaseError,
                InterfaceError,
                Error,
        ):
            db_exc_type = getattr(self.wrapper.Database, dj_exc_type.__name__)
            if issubclass(exc_type, db_exc_type):
                dj_exc_value = dj_exc_type(*exc_value.args)
                # Only set the 'errors_occurred' flag for errors that may make
                # the connection unusable.
                if dj_exc_type not in (DataError, IntegrityError):
                    self.wrapper.errors_occurred = True
                raise dj_exc_value.with_traceback(traceback) from exc_value

    def __call__(self, func):
        # Note that we are intentionally not using @wraps here for performance
        # reasons. Refs #21109.
        def inner(*args, **kwargs):
            with self:
                return func(*args, **kwargs)
        return inner

总结

这篇文章深入探讨了 Python 上下文管理器的应用与实现。首先，我们对比了在没有使用 with 语句和使用 with 语句时操作文件的代码差异，明确了 with 语句的优势，使代码结构更为简洁。接着，我们深入研究了 with 语句的实现原理，强调只需实现 __enter__ 和 __exit__ 方法的实例，即可借助 with 语法块高效管理资源。

然后我们也简单了解了下异步上下文管理器和async with以管理资源。

进一步介绍了 Python 标准库中 contextlib 模块，它提供了更便捷的上下文管理器实现方式。我们学习了如何使用 contextmanager 、ContextDecorator装饰器和 closing 等方法来优雅地处理资源。最后，通过多个实际例子展示了上下文管理器在管理资源和执行前后逻辑方面的价值。

综上所述，通过在开发中广泛运用上下文管理器，我们能够显著提升代码结构和可维护性。因此，强烈推荐在项目中积极采用这一优秀的编程模式。

参考

《流畅的python》第一版及第二版
《深入理解python特性》
docs.python.org/3/library/c…
realpython.com/python-with…
medium.com/@praba23089…
book.pythontips.com/en/latest/c…