Python的上下文管理器

编写自定义上下文管理器
从生成器到上下文管理器
将上下文管理器编写为装饰器
嵌套上下文
组合多个上下文管理器
使用上下文管理器创建SQLAlchemy会话
使用上下文管理器进行抽象处理异常处理
使用上下文管理器跨Http请求的持久参数
备注
资源资源
Python的上下文管理器非常适合资源管理和阻止泄漏抽象的传播。您可能在打开文件或数据库连接时使用了它。通常，它以这样的with语句开头：
1
2
with open("file.txt", "wt") as f:
f.write("contents go here")

在上述情况下，file.txt当执行流程超出范围时，将自动关闭。这等效于编写：

try:
    f = open("file.txt", "wt")
    text = f.write("contents go here")
finally:
    f.close()

编写自定义上下文管理器

要编写自定义上下文管理器，您需要创建一个包含 enter和exit方法的类。让我们重新创建一个自定义上下文管理器，它将执行与上述相同的工作流程。

class CustomFileOpen:
    """Custom context manager for opening files."""

    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.f = open(self.filename, self.mode)
        return self.f

    def __exit__(self, *args):
        self.f.close()

您可以像常规上下文管理器一样使用上述类。

1 2	with CustomFileOpen("file.txt", "wt") as f: f.write("contents go here")

关于上下文管理器

通过使用enter和exit方法编写类来创建上下文管理器并不困难。但是，通过使用contextlib.contextmanager装饰器定义它们，可以达到更好的简洁性。该装饰器将生成器功能转换为上下文管理器。创建上下文管理器装饰器的蓝图如下所示：

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

当将上下文管理器与with语句一起使用：

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

相当于：

<setup>
try:
    <variable> = <value>
    <body>
finally:
    <cleanup>

设置代码位于try..finally块之前。注意发电机屈服的点。这是嵌套在with语句中的代码块的执行位置。代码块完成后，将恢复生成器。如果该块中发生未处理的异常，则在yield发生该事件的位置在生成器内部重新引发该异常，然后finally执行该块。如果未发生未处理的异常，则代码会正常进行到finally运行清除代码的块。

让我们CustomFileOpen用contextmanager装饰器实现相同的上下文管理器。

from contextlib import contextmanager


@contextmanager
def CustomFileOpen(filename, method):
    """Custom context manager for opening a file."""

    f = open(filename, method)
    try:
        yield f

    finally:
        f.close()

现在使用：

1 2	with CustomFileOpen("file.txt", "wt") as f: f.write("contents go here")

将上下文管理器编写为装饰器

您还可以将上下文管理器用作装饰器。为此，在定义类时，您必须从contextlib.ContextDecorator类继承。让我们制作一个RunTime装饰器，该装饰器将应用于文件打开功能。装饰者将：

打印用户提供的功能说明
打印运行该功能所需的时间

from contextlib import ContextDecorator
from time import time


class RunTime(ContextDecorator):
    """Timing decorator."""

    def __init__(self, description):
        self.description = description

    def __enter__(self):
        print(self.description)
        self.start_time = time()

    def __exit__(self, *args):
        self.end_time = time()
        run_time = self.end_time - self.start_time
        print(f"The function took {run_time} seconds to run.")

例如：

@RunTime("This function opens a file")
def custom_file_write(filename, mode, content):
    with open(filename, mode) as f:
        f.write(content)

print(custom_file_write("file.txt", "wt", "jello"))
This function opens a file
The function took 0.0005390644073486328 seconds to run.
None

您还可以通过contextlib.contextmanager装饰器创建相同的装饰器。

from contextlib import contextmanager


@contextmanager
def runtime(description):

    print(description)
    start_time = time()
    try:
        yield
    finally:
        end_time = time()
        run_time = end_time - start_time
        print(f"The function took {run_time} seconds to run.")

嵌套上下文

您可以嵌套多个上下文管理器以同时管理资源。考虑以下虚拟管理器：

from contextlib import contextmanager


@contextmanager
def get_state(name):
    print("entering:", name)
    yield name
    print("exiting :", name)


# multiple get_state can be nested like this
with get_state("A") as A, get_state("B") as B, get_state("C") as C:
    print("inside with statement:", A, B, C)
entering: A
entering: B
entering: C
inside with statement: A B C
exiting : C
exiting : B
exiting : A

请注意它们的顺序。上下文管理器视为堆栈，应以与输入相反的顺序退出。如果发生异常，则此顺序很重要，因为任何上下文管理器都可以抑制该异常，此时其余的管理器甚至都不会收到此通知。exit还允许该方法引发其他异常，然后其他上下文管理器应能够处理该新异常。

组合多个上下文管理器

您也可以组合多个上下文管理器。让我们考虑这两个经理。

from contextlib import contextmanager


@contextmanager
def a(name):
    print("entering a:", name)
    yield name
    print("exiting a:", name)


@contextmanager
def b(name):
    print("entering b:", name)
    yield name
    print("exiting b:", name)

现在，使用装饰器语法将两者结合起来。以下函数采用上述定义管理器a，b并返回组合的上下文管理器ab。

@contextmanager
def ab(a, b):
    with a("A") as A, b("B") as B:
        yield (A, B)

with ab(a, b) as AB:
    print("Inside the composite context manager:", AB)
entering a: A
entering b: B
Inside the composite context manager: ('A', 'B')
exiting b: B
exiting a: A

如果您有可变数量的上下文管理器，并且想要适当地组合它们，contextlib.ExitStack可以在这里提供帮助。让我们ab使用重写上下文管理器ExitStack。此函数将各个上下文管理器及其参数作为元组，并返回组合的管理器。

from contextlib import contextmanager, ExitStack


@contextmanager
def ab(cms, args):
    with ExitStack() as stack:
        yield [stack.enter_context(cm(arg)) for cm, arg in zip(cms, args)]
with ab((a, b), ("A", "B")) as AB:
    print("Inside the composite context manager:", AB)
entering a: A
entering b: B
Inside the composite context manager: ['A', 'B']
exiting b: B
exiting a: A
ExitStack如果您要优雅地管理多个资源，也可以使用它。例如，假设您需要根据目录中多个文件的内容创建一个列表。让我们看看，如何通过强大的资源管理来避免意外的内存泄漏。

from contextlib import ExitStack
from pathlib import Path

# ExitStack ensures all files are properly closed after o/p
with ExitStack() as stack:
    streams = (
        stack.enter_context(open(fname, "r")) for fname in Path("src").rglob("*.py")
    )
    contents = [f.read() for f in streams]

使用上下文管理器创建SQLAlchemy会话

如果您熟悉SQLALchemy，Python的SQL工具包和对象关系映射器，那么您可能知道的用法Session来运行查询。Session基本上，A 将所有查询转换为事务并使其原子化。上下文管理器可以帮助您以非常优雅的方式编写事务会话。SQLAlchemy中的基本查询工作流程可能如下所示：

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from contextlib import contextmanager

# an Engine, which the Session will use for connection resources
some_engine = create_engine("sqlite://")

# create a configured "Session" class
Session = sessionmaker(bind=some_engine)


@contextmanager
def session_scope():
    """Provide a transactional scope around a series of operations."""
    session = Session()
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()

上面的摘录使用上下文管理器创建了一个内存中SQLite连接和一个session_scope函数。session_scope函数负责自动处理异常情况下的提交和回滚。该session_scope函数可以通过以下方式用于运行查询：

1
2
3

with session_scope() as session:
    myobject = MyObject("foo", "bar")
    session.add(myobject)

使用上下文管理器进行抽象并进行异常处理

这是我绝对喜欢的上下文管理器用例。假设您要编写一个函数，但希望避免异常处理逻辑。具有复杂日志记录的异常处理逻辑通常会混淆函数的核心逻辑。您可以编写一个装饰器类型的上下文管理器，该管理器将为您处理异常并将这些附加代码与主逻辑解耦。让我们写一个装饰，将处理ZeroDivisionError和TypeError同步。

from contextlib import contextmanager


@contextmanager
def errhandler():
    try:
        yield
    except ZeroDivisionError:
        print("This is a custom ZeroDivisionError message.")
        raise
    except TypeError:
        print("This is a custom TypeError message.")
        raise

现在，在发生这些异常的函数中使用它。

@errhandler()
def div(a, b):
    return a // b
div("b", 0)
This is a custom TypeError message.
---------------------------------------------------------------------------

TypeError                                 Traceback (most recent call last)

<ipython-input-43-65497ed57253> in <module>
----> 1 div('b',0)

/usr/lib/python3.8/contextlib.py in inner(*args, **kwds)
        73         def inner(*args, **kwds):
        74             with self._recreate_cm():
---> 75                 return func(*args, **kwds)
        76         return inner
        77


<ipython-input-42-b7041bcaa9e6> in div(a, b)
        1 @errhandler()
        2 def div(a, b):
----> 3     return a // b


TypeError: unsupported operand type(s) for //: 'str' and 'int'

您会看到errhandler装饰器正在为您完成繁重的工作。很整洁吧？

以下是使用上下文管理器将错误处理异常与主逻辑解耦的更复杂的示例。它还从main方法中隐藏了详细的日志记录逻辑。

import logging
from contextlib import contextmanager
import traceback
import sys

logging.getLogger(__name__)

logging.basicConfig(
    level=logging.INFO,
    format="\n(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.FileHandler("./debug.log"), logging.StreamHandler()],
)


class Calculation:
    """Dummy class for demonstrating exception decoupling with contextmanager."""

    def __init__(self, a, b):
        self.a = a
        self.b = b

    @contextmanager
    def errorhandler(self):
        try:
            yield
        except ZeroDivisionError:
            print(
                f"Custom handling of Zero Division Error! Printing "
                "only 2 levels of traceback.."
            )
            logging.exception("ZeroDivisionError")

    def main_func(self):
        """Function that we want to save from nasty error handling logic."""

        with self.errorhandler():
            return self.a / self.b


obj = Calculation(2, 0)
print(obj.main_func())

这将返回

(asctime)s [ERROR] ZeroDivisionError
Traceback (most recent call last):
    File "<ipython-input-44-ff609edb5d6e>", line 25, in errorhandler
    yield
    File "<ipython-input-44-ff609edb5d6e>", line 37, in main_func
    return self.a / self.b
ZeroDivisionError: division by zero


Custom handling of Zero Division Error! Printing only 2 levels of traceback..
None

使用上下文管理器跨Http请求的持久参数

上下文管理器的另一个很好的用例是使参数在多个http请求中保持不变。Python的requests库有一个Session对象，可让您轻松实现这一目标。因此，如果您要向同一主机发出多个请求，则基础TCP连接将被重用，这可以显着提高性能。以下示例直接取自请求的官方文档。让我们在请求中保留一些cookie。

with requests.Session() as session:
    session.get("http://httpbin.org/cookies/set/sessioncookie/123456789")
    response = session.get("http://httpbin.org/cookies")
    print(response.text)

结果为：

{
  "cookies": {
    "sessioncookie": "123456789"
  }
}

最后说明
所有代码片段均已针对python更新3.8。为避免冗余，我故意排除了带语句嵌套的示例，现在不建议使用contextlib.nested函数来创建嵌套的上下文管理器。