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

根据对 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()

Answer 2

如果您在 timer class 中还重写了从 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

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

Get function name when ContextDecorator is used as a decorator

python

decorator

python-decorators