python - 带参数和不带参数的上下文装饰器

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

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

诀窍是还将上下文装饰器的元类设为 ContextDecorator本身，然后覆盖 __call__方法来检查它是否正在传递标签(正常情况)或函数(无 parent 情况)。

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>