使用和不使用参数的上下文装饰器

Question

我想将上下文装饰器与使用或不使用参数的可能性结合起来。

让我们从一个可以带参数和不带参数的装饰器开始，例如：

import functools


def decorator(func=None, *, label=""):
    if func is None:
        return functools.partial(decorator, label=label)

    @functools.wraps(func)
    def wrap(*args, **kwargs):
        result = func(*args, **kwargs)
        print(f"RESULT {label}: {result}")
        return result

    return wrap


if __name__ == "__main__":

    @decorator(label="with arguments")
    def dec_args():
        return 1

    @decorator
    def dec_no_args():
        return 0

    dec_args()
    dec_no_args()

还有ContextDecorator 可以用作上下文管理器或装饰器：

from contextlib import ContextDecorator

class ctxtdec(ContextDecorator):
    def __init__(self, label:str=""):
        self.label = label
        print(f"initialized {self.label}")

    def __enter__(self):
        print(f"entered {self.label}")

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"exited {self.label}")

if __name__ == "__main__":
    def testfunc():
        for n in range(10 ** 7):
            n ** 0.5

    @ctxtdec("decorated")
    def decorated():
        testfunc()

    with ctxtdec("square rooting"):
        testfunc()
    decorated()

但我也希望它也能正常工作：

    @ctxtdec
    def decorated():
        testfunc()

Answer 1

警告：它并不漂亮，而且我永远不会真正使用它，但我很好奇，所以我让它工作了。可能有人也可以清理一下。

诀窍是让你的上下文装饰器的 metaclass 本身成为 ContextDecorator，然后重写 __call__ 方法来检查它是否被传递了标签（正常情况）或功能（无父母情况）。

from contextlib import ContextDecorator

class CtxMeta(type, ContextDecorator):
    def __enter__(self):
        print(f"entered <meta-with>")

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"exited <meta-with>")

    def __call__(cls, func_or_label=None, *args, **kwds):
        if callable(func_or_label):
            return type.__call__(cls, "<meta-deco>", *args, **kwds)(func_or_label)
        return type.__call__(cls, func_or_label, *args, **kwds)

然后，您的原始装饰器类与以前相同，但添加了元类声明：

class ctxtdec(ContextDecorator, metaclass=CtxMeta):
    def __init__(self, label:str=""):
        self.label = label
        print(f"initialized {self.label}")

    def __enter__(self):
        print(f"entered {self.label}")

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"exited {self.label}")

现在我们可以同时测试它（作为装饰器或上下文管理器）：

if __name__ == "__main__":
    def testfunc():
        for n in range(10 ** 7):
            n ** 0.5

    @ctxtdec("decorated")
    def decorated():
        testfunc()
    decorated()

    with ctxtdec("square rooting"):
        testfunc()

    @ctxtdec
    def deco2():
        testfunc()    
    deco2()

    with ctxtdec:
        testfunc()

还有输出：

initialized decorated
entered decorated
exited decorated
initialized square rooting
entered square rooting
exited square rooting
initialized <meta-deco>
entered <meta-deco>
exited <meta-deco>
entered <meta-with>
exited <meta-with>