使用和不使用参数的上下文修饰符

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

Caveat：这并不漂亮，我永远不会真正使用它，但是我很好奇，所以我使它起作用。也许有人也可以多清理一点。

诀窍是还使上下文修饰符的metaclass成为ContextDecorator本身，然后重写__call__方法以检查它是否传递了标签（正常情况）或函数（paren-较少的情况）。

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>