Context decorator that works with and without arguments

Question

I would like to combine a context decorator with the possiblity to work with or without arguments.

Lets start with a decorator that works both with and without arguments, for example:

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()

And there is the ContextDecorator which can be used as a contextmanager or a decorator:

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()

But I would also like this to work as well:

    @ctxtdec
    def decorated():
        testfunc()

Answer 1

Caveat: It's not pretty, and I would never actually use this, but I was curious so I made it work. Might be someone can clean it up a bit more too.

The trick is to also make the metaclass of your context decorator a ContextDecorator itself, and then override the __call__ method to check whether it's being passed a label (normal situation) or a function (paren-less situation).

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)

Then, your original decorator class stays the same as before, but with the addition of a metaclass declaration:

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}")

And now we can test it both ways (as a decorator or a context-manager):

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()

And the output:

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>