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?
If you also override the __call__() method inherited from the ContextDecorator base class in your timer 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
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()