Python上下文管理器的魔力

点击关注我哦

一篇文章带你了解Python上下文管理器的魔力

小编将为您准备一份很棒的Python上下文管理器使用指南，这将使您的代码更具可读性和可靠性，降低您的错误发生率。

资源管理器是我们在任何编程语言中都需要使用的工具之一。无论是处理锁、文件、会话还是数据库链接，我们都必须确保关闭并释放这些资源，以确保他们能正确运行。通常，可以使用try / finally来做到这一点：在try代码块中使用资源，然后在finally代码块中进行处理。 但是在Python中有更好的方法，那使用with语句实现上下文管理协议。

因此在本文中，我们将探讨它的含义和工作原理，最重要的是我们会知道在哪里能够找到以及如何实现自己的出色上下文管理器！

什么是上下文管理器呢？

即使您从未听说过Python的上下文管理器，但根据上面的介绍，也已经知道它可以替代try / finally模块的作用了。 它是使用打开文件时常用的with语句实现的。 与try / finally一样，引入此模式是为了确保即使在发生异常或程序终止的情况下，也将在块的末尾执行某些操作。

从表面上看，上下文管理协议只是使用with语句的代码块。实际上，它由2种特殊的（dunder）方法-__enter__和__exit__组成，它们分别用于设置和拆解。

当代码中遇到with语句时，将触发__enter__方法，并将其返回值放入as限定符之后的变量中。 在with块的主体执行之后，将调用__exit__方法来执行拆解-履行finally块的作用。

# Using try/finally
import time


start = time.perf_counter()  # Setup
try:  # Actual body
    time.sleep(3)
finally:  # Teardown
    end = time.perf_counter()
    elapsed = end - start
    
print(elapsed)


# Using Context Manager
with Timer() as t:
    time.sleep(3)


print(t.elapsed)

上面的代码同时显示了使用try / finally的版本和使用with语句实现简单计时器的更优雅的版本。 我们在上面提到，实现此类上下文管理器需要__enter__和__exit__，但是我们将如何创建它们呢？ 让我们看一下这个Timer类的代码：

# Implementation of above context manager
class Timer:
    def __init__(self):
        self._start = None
        self.elapsed = 0.0


    def start(self):
        if self._start is not None:
            raise RuntimeError('Timer already started...')
        self._start = time.perf_counter()


    def stop(self):
        if self._start is None:
            raise RuntimeError('Timer not yet started...')
        end = time.perf_counter()
        self.elapsed += end - self._start
        self._start = None


    def __enter__(self):  # Setup
        self.start()
        return self


    def __exit__(self, *args):  # Teardown
        self.stop()

此代码段显示了同时实现__enter__和__exit__方法的Timer类。 __enter__方法仅启动计时器并返回self。with语句主体完成后，将使用3个参数（异常类型，异常值和回溯）调用__exit__方法。 如果一切都很好，with语句将全部等于None。 如果引发异常，则将使用异常数据填充这些数据，我们可以在__exit__方法中进行处理。 在这种情况下，我们忽略异常处理，只停止计时器并计算经过的时间，并将其存储在上下文管理器的属性中。

我们已经在这里看到了with语句的实现和示例用法，但是为了更直观地了解实际情况，让我们看一下如何在没有Python's syntax sugar的情况下调用这些特殊方法：

manager = Timer()
manager.__enter__()  # Setup
time.sleep(3)  # Body
manager.__exit__(None, None, None)  # Teardown
print(manager.elapsed)

既然我们已经确定了上下文管理器的含义、它的工作方式和实现方式，为了有更多的动力从try / finally转换为with语句，让我们看一下使用它的好处。

第一个好处是，整个设置和拆卸都在上下文管理器对象的控制下进行。 这样可以防止错误并减少样板代码，从而使API更安全，更易于使用。 使用它的另一个原因是with块突出显示了关键部分，并鼓励您减少该部分中的代码量，这通常也是一种好习惯。 最后-最后但并非最不重要-这是一个很好的重构工具，它可以排除常见的设置和拆卸代码并将其移至__enter__和__exit__方法的单个位置。

话虽如此，如果您以前没有使用过，我希望我能说服您开始使用上下文管理器而不是try/finally。 现在，让我们来看一些很酷且有用的上下文管理器，您应该开始将它们包括在代码中！

使用@contextmanager使一切变得简单

在上一节中，我们探讨了如何使用__enter__和__exit__方法实现上下文管理器。 这很简单，但是我们可以使用contextlib，更具体地说使用@contextmanager使其更加简单。

@contextmanager是一个装饰器，可用于编写独立的上下文管理功能。 因此，我们不需要创建整个类并实现__enter__和__exit__方法，而是要做的就是创建单个生成器：

from contextlib import contextmanager
from time import time, sleep


@contextmanager
def timed(label):
    start = time()  # Setup - __enter__
    print(f"{label}: Start at {start}")
    try:  
        yield  # yield to body of `with` statement
    finally:  # Teardown - __exit__
        end = time()
        print(f"{label}: End at {end} ({end - start} elapsed)")


with timed("Counter"):
    sleep(3)


# Counter: Start at 1599153092.4826472
# Counter: End at 1599153095.4854734 (3.00282621383667 elapsed)

该片段实现了与上一节中的Timer类非常相似的上下文管理器。 但是这次，我们需要的代码更少了。 这小段代码分为两部分-屈服之前的一切和屈服之后的一切。yield之前的代码执行__enter__方法的工作，而yield本身就是__enter__方法的return语句。yield之后的所有内容都是__exit__方法的一部分。

如您在上面看到的，使用这样的单个函数创建上下文管理器需要使用try / finally，因为如果with语句的主体中发生异常，则它将与yield一起引发，并且我们最终需要处理它 对应于__exit__方法的块。

正如我已经提到的，它可以用于独立的上下文管理器。 但是，它不适用于需要成为对象一部分的上下文管理器，例如连接或锁定。

尽管使用单个功能构建上下文管理器会迫使您使用try / finally，并且只能与更简单的用例一起使用，但在我看来，对于构建更精简的上下文管理器，它仍然是一种优雅而实用的选择。

现实生活中的例子

现在，让我们从理论转向实践，自己构建实用的上下文管理器。

记录上下文管理器

当需要尝试查找代码中的某些错误时，您可能首先会在日志中查找问题的根本原因。 但是，这些日志可能默认情况下设置为错误或警告级别，可能不足以进行调试。 更改整个程序的日志级别应该很容易，但是针对代码的特定部分更改日志级别可能会更复杂-但是，可以使用以下上下文管理器轻松解决：

import logging
from contextlib import contextmanager


@contextmanager
def log(level):
    logger = logging.getLogger()
    current_level = logger.getEffectiveLevel()
    logger.setLevel(level)
    try:
        yield
    finally:
        logger.setLevel(current_level)


def some_function():
    logging.debug("Some debug level information...")
    logging.error('Serious error...')
    logging.warning('Some warning message...')


with log(logging.DEBUG):
    some_function()


# DEBUG:root:Some debug level information...
# ERROR:root:Serious error...
# WARNING:root:Some warning message...

超时上下文管理器

在本文的开头，我们讨论了代码的计时块。 我们将在这里尝试将超时设置为with语句所包围的块：

import signal
from time import sleep


class timeout:
    def __init__(self, seconds, *, timeout_message=""):
        self.seconds = int(seconds)
        self.timeout_message = timeout_message


    def _timeout_handler(self, signum, frame):
        raise TimeoutError(self.timeout_message)


    def __enter__(self):
        signal.signal(signal.SIGALRM, self._timeout_handler)  # Set handler for SIGALRM
        signal.alarm(self.seconds)  # start countdown for SIGALRM to be raised


    def __exit__(self, exc_type, exc_val, exc_tb):
        signal.alarm(0)  # Cancel SIGALRM if it's scheduled
        return exc_type is TimeoutError  # Suppress TimeoutError




with timeout(3):
    # Some long running task...
    sleep(10)

上面的代码为此上下文管理器声明了一个称为timeout的类，因为该任务无法在单个函数中完成。 为了实现这种超时，我们还需要使用信号-更具体地说是SIGALRM。 我们首先使用signal.signal（...）将处理程序设置为SIGALRM，这意味着当SIGALRM由内核引发时，将调用处理程序函数。 对于此处理函数（_timeout_handler），它所做的只是引发TimeoutError，如果未及时完成，它将在with语句的主体中停止执行。 有了处理程序之后，我们还需要以指定的秒数开始倒数，这是通过signal.alarm（self.seconds）完成的。

至于__exit__方法-如果上下文管理器主体设法在时间到期之前完成，则SIGALRM将被signal.alarm（0）取消，程序可以继续。 另一方面-如果由于超时而引发信号，则_timeout_handler将引发TimeoutError，它将被__exit__捕获并抑制，with语句的主体将被中断，其余代码可以继续执行。

使用已经存在的内容

除了上述上下文管理器外，标准库或其他常用库（例如request或sqlite3）中已经有很多有用的库。 因此，让我们看看在这里可以找到什么。

临时更改小数精度

如果您要进行大量的数学运算并需要特定的精度，那么您可能会遇到以下情况：可能希望暂时更改十进制数字的精度：

from decimal import getcontext, Decimal, setcontext, localcontext, Context


# Bad
old_context = getcontext().copy()
getcontext().prec = 40
print(Decimal(22) / Decimal(7))
setcontext(old_context)


# Good
with localcontext(Context(prec=50)):
    print(Decimal(22) / Decimal(7))  # 3.1428571428571428571428571428571428571428571428571


print(Decimal(22) / Decimal(7))      # 3.142857142857142857142857143

上面的代码演示了不带上下文管理器和带上下文管理器的两种选择。 第二种选择显然更短并且更具可读性。 它还将临时上下文排除在外，这使得它不太容易出错。

contextlib

在使用@contextmanager时，我们已经窥探了contextlib，但是还有更多可以使用的东西。作为第一个示例，让我们看一下redirect_stdout和redirect_stderr：

import sys
from contextlib import redirect_stdout


# Bad
with open("help.txt", "w") as file:
    stdout = sys.stdout
    sys.stdout = file
    try:
        help(int)
    finally:
        sys.stdout = stdout


# Good
with open("help.txt", "w") as file:
    with redirect_stdout(file):
        help(int)

如果您具有默认情况下将所有内容输出到stdout或stderr的工具或功能，但是您希望它在其他位置输出数据-例如 归档-那么这两个上下文管理器可能会很有帮助。 与前面的示例一样，这大大提高了代码的可读性并消除了不必要的视觉噪声。

contextlib中的另一个方便使用的是抑制上下文管理器，它将抑制所有不需要的异常和错误：

import os
from contextlib import suppress


try:
    os.remove('file.txt')
except FileNotFoundError:
    pass




with suppress(FileNotFoundError):
    os.remove('file.txt')

正确处理异常绝对是可取的，但是有时您只需要摆脱讨厌的DeprecationWarning，并且此上下文管理器至少会使它易于阅读。

 

下面提到的是contextlib的最后一个实际上也是我的最喜欢特性，它称为close：

# Bad
try:
    page = urlopen(url)
    ...
finally:
    page.close()


# Good
from contextlib import closing


with closing(urlopen(url)) as page:
    ...

上下文管理器，用于更好的测试

如果您希望人们使用，阅读或维护您编写的测试，则必须使他们可读且易于理解，那么mock.patch上下文管理器可以帮助您：

# Bad
import requests
from unittest import mock
from unittest.mock import Mock


r = Mock()
p = mock.patch('requests.get', return_value=r)
mock_func = p.start()
requests.get(...)
# ... do some asserts
p.stop()


# Good
r = Mock()
with mock.patch('requests.get', return_value=r):
    requests.get(...)
    # ... do some asserts

在上下文管理器中使用mock.patch可以使您摆脱不必要的.start（）和.stop（）调用，并帮助您定义此特定mock的明确范围。 关于这一点的一件好事是，即使它是标准库（因此也是unittest）的一部分，它也可以与unittest和pytest一起使用。

在谈到pytest时，我们还要展示至少一个非常有用的上下文管理器：

import pytest, os


with pytest.raises(FileNotFoundError, message="Expecting FileNotFoundError"):
    os.remove('file.txt')

这个例子展示了pytest.raises的非常简单的用法，它断言代码块引发了所提供的异常。 如果不是，则测试失败。 这对于测试预期会引发异常或失败的代码路径非常方便。

跨请求保留会话

现在从pytest转到另一个很棒的库--requests。 通常，您可能需要在HTTP请求之间保留cookie，需要保持TCP连接保持活动状态，或者只是想对同一主机执行多个请求。requests提供了不错的上下文管理器来帮助应对这些挑战-即-用于管理会话：

import requests


with requests.Session() as session:
    session.request(method=method, url=url, **kwargs)

除了解决上述问题之外，此上下文管理器还可以提高性能，因为它将重用基础连接，因此避免为每个请求/响应对打开新连接。

管理SQLite事务

最后但并非最不重要的一点是，还有用于管理SQLite事务的上下文管理器。 除了使代码更简洁外，此上下文管理器还提供在异常情况下回滚更改的功能，以及在with语句主体成功完成时自动提交的功能：

import sqlite3
from contextlib import closing


# Bad
connection = sqlite3.connect(":memory:")
try:
    connection.execute("INSERT INTO employee(firstname, lastname) values (?, ?)", ("John", "Smith",))
except sqlite3.IntegrityError:
    ...


connection.close()


# Good
with closing(sqlite3.connect(":memory:")) as connection:
    with connection:
        connection.execute("INSERT INTO employee(firstname, lastname) values (?, ?)", ("John", "Smith",))

在此示例中，您还可以看到上下文管理器中closing的使用，它有助于处理不再使用的连接对象，这进一步简化了此代码，并确保我们不会挂起任何连接。

结 论

我要强调的一件事是，上下文管理器不仅是资源管理工具，而且还允许您提取和分解任意一对操作的通用设置和拆卸的功能，而不仅仅是诸如锁或网络连接之类的通用用例。  它也是Python的强大功能之一，几乎不可能以其他任何语言找到。 整洁优雅，因此希望本文向您展示了上下文管理器的功能，并向您介绍了在代码中使用它们的更多方法。

·  END  ·

HAPPY LIFE