ContextDecorator用作装饰器时获取函数名

Question

我有以下上下文管理器和装饰器来为任何给定的函数或代码块计时：

import time
from contextlib import ContextDecorator


class timer(ContextDecorator):
    def __init__(self, label: str):
        self.label = label

    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

您可以将其用作上下文管理器：

with timer("my code block"):
    time.sleep(2)

# my code block took 2.0 seconds

您也可以将其用作装饰器：

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

my_func()

# my_func took 3.0 seconds

我唯一不喜欢的是在用作装饰器时必须手动将函数名称作为label 传递。如果没有传递标签，我希望装饰器自动使用函数名称：

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

my_func()

# my_func took 3.0 seconds

有什么办法吗？

Answer 1

如果您还覆盖了从 timer 类中的 ContextDecorator base class 继承的 __call__() 方法，并为 label 参数的初始化程序添加一个唯一的默认值，您可以检查它并调用函数时获取函数的__name__：

import time
from contextlib import ContextDecorator


class timer(ContextDecorator):
    def __init__(self, label: str=None):
        self.label = label

    def __call__(self, func):
        if self.label is None:  # Label was not provided
            self.label = func.__name__  # Use function's 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


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

my_func()  # -> my_func took 3.0 seconds

Answer 2

根据对ContextDecorator 的source 的检查，似乎无法将包装函数的名称提供给上下文管理器。相反，您可以创建自己的ContextDecorator 版本，覆盖__call__

import time
import functools, contextlib
class _ContextDecorator(contextlib.ContextDecorator):
   def __call__(self, func):
      self.f_name = func.__name__
      @functools.wraps(func)
      def wrapper(*args, **kwargs):
          with self._recreate_cm():
             return func(*args, **kwargs)
      return wrapper

用法：

class timer(_ContextDecorator):
   def __init__(self, label: str = None):
      self.f_name = label
   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.f_name} took {net_time:.1f} seconds")
      return False

with timer('my_func'):
   time.sleep(2)

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

my_func()