Get function name when ContextDecorator is used as a decorator

Question

I have the following context manager and decorator to time any given function or code block:

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

You can use it as a context manager:

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

# my code block took 2.0 seconds

You can also use it as a decorator:

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

my_func()

# my_func took 3.0 seconds

The only thing I don't like is having to manually pass the function name as the label when it's used as a decorator. I would love for the decorator to automatically use the function name if no label is passed:

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

my_func()

# my_func took 3.0 seconds

Is there any way to do this?

Answer 1

If you also override the __call__() method inherited from ContextDecorator base class in your class, and add a unique default value to the initializer for the label argument, you can check for that and grab the function's __name__ when it's called:

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

Based on an examination of ContextDecorator 's source , there does not seem to be any way the name of the wrapped function can be made available to the contextmanager. Instead, you can create your own version of ContextDecorator , overriding __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

Usage:

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