Python 中的条件上下文管理器和装饰器

Question

我的代码库中散布着以下 Python 代码，如果设置了 SHOULD_PROFILE 环境变量，它会分析代码块并将结果发送到监控解决方案。

from contextlib import ExitStack

with ExitStack() as stack:
    if os.environ.get("SHOULD_PROFILE"):
        stack.enter_context(profile())
    ...

我想将此代码段合并为一个上下文管理器/装饰器，以便它可以用作上下文管理器和装饰器：

with profile_if_enabled():
    ...

# AND

@profile_if_enabled()
def func():
    ....

这是我想出的方法，但它不起作用。

from contextlib import ContextDecorator, ExitStack

class profile_if_enabled(ContextDecorator):
    def __enter__(self):
        with ExitStack() as stack:
            if os.environ.get("SHOULD_PROFILE"):
                stack.enter_context(profile())
            return self

    def __exit__(self, *exc):
        return False

知道我做错了什么吗？

Answer 1

您当前的尝试失败了，因为一旦您 return self，您就离开了 with ExitStack() as stack: 块，并且 ExitStack 执行清理。您需要在 __exit__ 中执行清理，而不是在 __enter__ returns.

中执行清理

ContextDecorator 需要一个可重用的、理想情况下可重入的上下文管理器，它不能很好地满足您的需求。即使你得到了一些工作，如果你试图从另一个配置文件函数调用一个配置文件函数，或者如果你尝试运行同时在不同线程中的两个配置文件函数，它很可能会中断。

手动编写装饰器会更干净:

def profile_if_enabled(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if os.environ.get('SHOULD_PROFILE'):
            with profile():
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)
    return wrapper

如果您确实想要同时用作装饰器和上下文管理器的东西，您可以检查是否存在参数：

def profile_if_enabled(func=None):
    if func is None:
        return profile()

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if os.environ.get('SHOULD_PROFILE'):
            with profile():
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)
    return wrapper

这将用作 with profile_if_enabled(): 或 @profile_if_enabled。请注意，与 ContextDecorator 不同，用作装饰器没有括号。

如果您不需要支持在执行期间更改 SHOULD_PROFILE 设置，您也可以重组代码以避免重复检查：

if os.environ.get('SHOULD_PROFILE'):
    def profile_if_enabled(func=None):
        if func is None:
            return profile()

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            with profile():
                return func(*args, **kwargs)
        return wrapper
else:
    def profile_if_enabled(func=None):
        if func is None:
            return contextlib.nullcontext()
        return func

Answer 2

使用 contextmanager 装饰器解决了这个问题：

import os
from contextlib import ExitStack, contextmanager

@contextmanager
def profile_if_enabled():
    with ExitStack() as stack:
        if os.environ.get("SHOULD_PROFILE"):
            stack.enter_context(profile())
        yield

Python 中的条件上下文管理器和装饰器

Conditional context manager and decorator in Python

python

decorator

contextmanager