上下文管理器和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 语句的使用场景：

1. 文件操作：open() 函数返回的文件对象就是一个上下文管理器，它负责文件的打开和关闭。with open("file.txt", "r") as file: 就是一个典型的使用案例。

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

1. 数据库连接：在数据库操作中，使用 with 语句可以确保在使用完数据库连接后正确地关闭连接，防止资源泄漏。

import sqlite3

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

1. 网络连接：类似于数据库连接，确保在网络连接使用完毕后关闭连接是良好的实践。

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)

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

1. 线程锁：threading 模块中的 Lock 对象可以作为上下文管理器，确保在使用完锁之后正确释放。

import threading

lock = threading.Lock()

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

1. 文件锁：在多进程环境下，可以使用文件锁确保多个进程对同一文件的互斥访问。

import fcntl

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

1. 时间测量：使用上下文管理器来测量代码块的执行时间。

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():
# 执行需要测量时间的代码

1. 测试资源管理：在测试中，可以使用上下文管理器确保在测试开始和结束时正确地分配和释放资源。

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: <class 'ZeroDivisionError'>
exc_value: division by zero
exc_tb: <traceback object at 0x7fa1d2adfe40>
__exit__
Traceback (most recent call last):
  File "/mnt/d/jinzhuan/rss/tpl_hf.py", line 15, in <module>
    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__ 方法都会被调用。

import time


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

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.end_time = time.time()
        elapsed_time = self.end_time - self.start_time
        print(f"Elapsed Time: {elapsed_time} seconds")


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 -> <!doctype html>
# https://realpython.com: type -> <!doctype html

使用asyc with语句

async with是with语句的异步版本，可以使用它来编写依赖于异步代码的上下文管理器。

import aiohttp
import asyncio


async def check(url):
    async with aiohttp.ClientSession() as session:
        async with session.get(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://realpython.com: type -> <!doctype html
# https://pycoders.com: status -> 200
# https://pycoders.com: type -> <!doctype html>

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

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

1. contextlib.contextmanager：

介绍： 这是一个装饰器，用于将一个生成器函数转化为上下文管理器。生成器的 yield 语句之前的代码块将在 __enter__ 方法中执行，而 yield 语句之后的代码块将在 __exit__ 方法中执行。

此装饰器本质： 我们先来看下这个装饰器的定义和帮助文档，其实讲的还是非常清楚的，巧妙的利用了生成器特性实现了上下文管理器。

def contextmanager(func):
  """@contextmanager decorator.

  Typical usage:

      @contextmanager
      def some_generator(<arguments>):
          <setup>
          try:
              yield <value>
          finally:
              <cleanup>

  This makes this:

      with some_generator(<arguments>) as <variable>:
          <body>

  equivalent to this:

      <setup>
      try:
          <variable> = <value>
          <body>
      finally:
          <cleanup>
  """
  @wraps(func)
  def helper(*args, **kwds):
      return _GeneratorContextManager(func, args, kwds)
  return helper

class _GeneratorContextManager(_GeneratorContextManagerBase,
                             AbstractContextManager,
                             ContextDecorator):
   """Helper for @contextmanager decorator."""

   def _recreate_cm(self):
       # _GCM instances are one-shot context managers, so the
       # CM must be recreated each time a decorated function is
       # called
       return self.__class__(self.func, self.args, self.kwds)

   def __enter__(self):
       # do not keep args and kwds alive unnecessarily
       # they are only needed for recreation, which is not possible anymore
       del self.args, self.kwds, self.func
       try:
           return next(self.gen)
       except StopIteration:
           raise RuntimeError("generator didn't yield") from None

   def __exit__(self, type, value, traceback):
       if type is None:
           try:
               next(self.gen)
           except StopIteration:
               return False
           else:
               raise RuntimeError("generator didn't stop")
       else:
           if value is None:
               # Need to force instantiation so we can reliably
               # tell if we get the same exception back
               value = type()
           try:
               self.gen.throw(type, value, traceback)
           except StopIteration as exc:
               # Suppress StopIteration *unless* it's the same exception that
               # was passed to throw().  This prevents a StopIteration
               # raised inside the "with" statement from being suppressed.
               return exc is not value
           except RuntimeError as exc:
               # Don't re-raise the passed in exception. (issue27122)
               if exc is value:
                   return False
               # Likewise, avoid suppressing if a StopIteration exception
               # was passed to throw() and later wrapped into a RuntimeError
               # (see PEP 479).
               if type is StopIteration and exc.__cause__ is value:
                   return False
               raise
           except:
               # only re-raise if it's *not* the exception that was
               # passed to throw(), because __exit__() must not raise
               # an exception unless __exit__() itself failed.  But throw()
               # has to raise the exception to signal propagation, so this
               # fixes the impedance mismatch between the throw() protocol
               # and the __exit__() protocol.
               #
               # This cannot use 'except BaseException as exc' (as in the
               # async implementation) to maintain compatibility with
               # Python 2, where old-style class exceptions are not caught
               # by 'except BaseException'.
               if sys.exc_info()[1] is value:
                   return False
               raise
           raise RuntimeError("generator didn't stop after throw()")

示例：

from contextlib import contextmanager


@contextmanager
def my_context():
    # Code before yield is __enter__
    print("Enter")
    yield
    # Code after yield is __exit__
    print("Exit")


with my_context():
    print("Inside the context")

yield后异常示例：

from contextlib import contextmanager


@contextmanager
def test():
    print('before')
    try:
        yield 'hello'
        # 这里发生异常必须自己处理异常逻辑否则不会向下执行
        a = 1 / 0
    finally:
        print('after')


with test() as t:
    print(t)
    
"""
before
hello
after
Traceback (most recent call last):
  File "/mnt/d/jinzhuan/rss/tpl_hf.py", line 16, in <module>
    print(t)
  File "/usr/lib/python3.8/contextlib.py", line 120, in __exit__
    next(self.gen)
  File "/mnt/d/jinzhuan/rss/tpl_hf.py", line 10, in test
    a = 1 / 0
ZeroDivisionError: division by zero
"""

1. contextlib.closing：

介绍： closing 是一个上下文管理器工具，用于确保在退出上下文时调用对象的close方法。通常，它用于处理那些没有实现上下文管理协议（__enter__和__exit__方法）的对象，但是有close方法的情况。

示例：

from contextlib import closing


class Test(object):

    # 定义了 close 方法才可以使用 closing 装饰器
    def close(self):
        print('closed')

    # with 块执行结束后 自动执行 close 方法


with closing(Test()):
    print('do something')

# Output:  
# do something  
# closed

从执行结果我们可以看到，with 语句块执行结束后，会自动调用 Test 实例的 close 方法。所以，对于需要自定义关闭资源的场景，我们可以使用这个方法配合 with 来完成。

1. contextlib.ExitStack：

介绍： ExitStack 是一个上下文管理器，它允许你动态地管理一组上下文管理器，无论这组管理器的数量是多少。它可以用于替代多个嵌套的 with 语句。

示例：

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')

1. contextlib.suppress：

介绍： suppress 是一个上下文管理器，用于忽略指定的异常。它允许您执行代码块，即使在存在异常的情况下也能继续执行。

示例：

from contextlib import suppress

with suppress(ZeroDivisionError):
    # 尝试执行可能抛出异常的代码
    result = 10 / 0

# 不会中断程序执行，即使发生了ZeroDivisionError
print("Continuing with the program")

1. contextlib.redirect_stdout 和 contextlib.redirect_stderr：

介绍： 这两个上下文管理器允许您在代码块内重定向标准输出和标准错误流。

示例：

from contextlib import redirect_stdout, redirect_stderr
import io

# 示例1: 将标准输出重定向到文件
with open('output.txt', 'w') as f:
    with redirect_stdout(f):
        print("This goes to the file instead of the console.")

# 示例2: 将标准错误重定向到字符串
error_output = io.StringIO()
with redirect_stderr(error_output):
    print("This goes to the error_output StringIO instead of the console.")

# 获取重定向的标准错误的内容
error_contents = error_output.getvalue()
print("Captured error output:", error_contents)

# This goes to the error_output StringIO instead of the console.
# Captured error output:

1. contextlib.ContextDecorator：

介绍： ContextDecorator是一个用于创建同时兼具上下文管理器和装饰器功能的基类。

示例：

from contextlib import ContextDecorator


class TestCm(ContextDecorator):
    a = 12

    def __enter__(self):
        print("enter")
        return self

    def __exit__(self, exc_type, exception, traceback):
        print("exit")
        TestCm.a = 12


@TestCm()
def play_testcm():
    TestCm.a = 8
    print("play_testcm, ", TestCm.a)


play_testcm()



# enter
# play_testcm,  8
# exit

1. contextlib.asynccontextmanager：

介绍： contextlib.asynccontextmanager是用于创建异步上下文管理器的装饰器。

示例：

import asyncio
from contextlib import asynccontextmanager


# 异步上下文管理器的例子
@asynccontextmanager
async def async_example():
    print("Enter async_example")

    # 在进入上下文时执行一些异步操作

    result = await asyncio.sleep(2)

    try:
        # 将资源（或状态）传递给上下文
        yield result
    finally:
        # 在退出上下文时执行一些异步操作
        print("Exit async_example")


# 使用异步上下文管理器
async def main():
    async with async_example() as result:
        # 在此处使用上下文中的资源
        print("Inside main:", result)


# 运行异步代码
asyncio.run(main())

# Enter async_example
# Inside main: None
# Exit async_example

这些工具使得使用上下文管理器更加方便和灵活。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分布式锁

import redis
from contextlib import contextmanager

# 假设已经创建了一个 Redis 连接
redis_client = redis.StrictRedis(host='localhost', port=6379, decode_responses=True)


@contextmanager
def lock(_redis_client, lock_key, expire):
    try:
        _locked = _redis_client.set(lock_key, 'locked', ex=expire, nx=True)
        yield _locked
    finally:
        _redis_client.delete(lock_key)


# 使用 lock 上下文管理器
with lock(redis_client, 'locked', 3) as locked:
    if not locked:
        print("Failed to acquire lock. Exiting...")
    else:
        print("Lock acquired. Performing some operation...")
        # 在这里执行需要锁的操作

# do something...

Redis事务和管道

import redis
from contextlib import contextmanager

# 假设已经创建了一个 Redis 连接
redis_client = redis.StrictRedis(host='localhost', port=6379, decode_responses=True)


@contextmanager
def pipeline(_redis_client):
    _pipe = _redis_client.pipeline()
    try:
        yield _pipe
        _pipe.execute()
    except Exception as exc:
        _pipe.reset()


# 使用 pipeline 上下文管理器
with pipeline(redis_client) as pipe:
    pipe.set('key1', 'a', ex=30)
    pipe.zadd('key2', {'a': 1})
    pipe.sadd('key3', 'a')

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

1. pipeline 函数使用 @contextmanager 装饰器，使得该函数返回一个上下文管理器。
2. 在 try 块内，创建了一个 Redis pipeline 对象 pipe。
3. yield pipe 将这个 pipeline 对象传递给 with 语句中的 as 子句，这样在 with 代码块中可以使用这个 pipeline 对象。
4. pipe.execute() 在 with 代码块结束时被调用，执行 Redis pipeline 中的所有命令。
5. except 块内，如果在 with 代码块中发生了异常，通过 pipe.reset() 重置 pipeline 对象，确保不会执行之前的命令。

异步文件上下文管理器

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…