继承ContextDecorator的装饰器如何自动获取被装饰函数的名称

核心思路 #

我们只需要修改timer类的初始化方法和__call__方法即可实现需求：

1. 将label参数改为可选，默认值为None
2. 利用ContextDecorator的__call__机制，在作为装饰器执行时如果检测到未传入label，自动读取被装饰函数的__name__属性作为标签
3. 可额外兼容省略括号的@timer装饰器写法，更符合使用习惯

修改后完整代码 #

import time
from contextlib import ContextDecorator
from typing import Optional, Callable, Union


class timer(ContextDecorator):
    # Python 3.10及以上版本可将参数注解改为label: Optional[Callable | str] = None
    def __init__(self, label: Optional[Union[Callable, str]] = None):
        # 兼容不带括号的装饰器用法：@timer
        if callable(label):
            self.func = label
            self.label = label.__name__
        else:
            self.label = label
            self.func = None

    def __call__(self, *args, **kwargs):
        # 不带括号的装饰器场景，直接执行被装饰函数
        if self.func is not None:
            with self:
                return self.func(*args, **kwargs)
        # 带括号的装饰器场景，读取函数名作为label
        func = args[0]
        if self.label is None:
            self.label = func.__name__
        return super().__call__(func)

    def __enter__(self):
        self.start_time = time.perf_counter()
        return self

    def __exit__(self, *exc):
        net_time = time.perf_counter() - self.start_time
        print(f"{self.label} took {net_time:.1f} seconds")
        return False

使用验证 #

1. 无参数装饰器用法（需求要求的效果） #

@timer()
def my_func():
    time.sleep(3)

my_func()
# 输出：my_func took 3.0 seconds

2. 省略括号的装饰器用法（额外兼容） #

@timer
def my_func2():
    time.sleep(1)

my_func2()
# 输出：my_func2 took 1.0 seconds

3. 自定义label的装饰器用法（原有能力保留） #

@timer("自定义函数标签")
def my_func3():
    time.sleep(2)

my_func3()
# 输出：自定义函数标签 took 2.0 seconds

4. 上下文管理器用法（原有能力保留） #

with timer("自定义代码块标签"):
    time.sleep(2)
# 输出：自定义代码块标签 took 2.0 seconds

内容的提问来源于stack exchange，提问作者Johnny Metz