# Imports
from collections.abc import Callable, Generator
from functools import wraps
from typing import Any, Literal, overload
from ..ctx.measure_time import MeasureTime
from ..print.message import progress
from .common import get_function_name, get_wrapper_name, set_wrapper_name
# Execution time decorator
# Regular function overloads (is_generator=False)
@overload
def measure_time[T](
func: Callable[..., T],
*,
printer: Callable[..., None] = progress,
message: str = "",
perf_counter: bool = True,
is_generator: Literal[False] = False
) -> Callable[..., T]: ...
@overload
def measure_time[T](
func: None = None,
*,
printer: Callable[..., None] = progress,
message: str = "",
perf_counter: bool = True,
is_generator: Literal[False] = False
) -> Callable[[Callable[..., T]], Callable[..., T]]: ...
# Generator function overloads (is_generator=True)
@overload
def measure_time[T](
func: Callable[..., Generator[T, None, None]],
*,
printer: Callable[..., None] = progress,
message: str = "",
perf_counter: bool = True,
is_generator: Literal[True]
) -> Callable[..., Generator[T, None, None]]: ...
@overload
def measure_time[T](
func: None,
*,
printer: Callable[..., None] = progress,
message: str = "",
perf_counter: bool = True,
is_generator: Literal[True]
) -> Callable[[Callable[..., Generator[T, None, None]]], Callable[..., Generator[T, None, None]]]: ...
[docs]
def measure_time[T](
func: Callable[..., T] | Callable[..., Generator[T, None, None]] | None = None,
*,
printer: Callable[..., None] = progress,
message: str = "",
perf_counter: bool = True,
is_generator: bool = False
) -> (
Callable[..., T]
| Callable[..., Generator[T, None, None]]
| Callable[[Callable[..., T]], Callable[..., T]]
| Callable[[Callable[..., Generator[T, None, None]]], Callable[..., Generator[T, None, None]]]
):
""" Decorator that will measure the execution time of a function and print it with the given print function
Args:
func (Callable[..., Any] | None): Function to decorate
printer (Callable): Function to use to print the execution time
(e.g. :py:func:`~stouputils.print.debug`, :py:func:`~stouputils.print.info`, :py:func:`~stouputils.print.warning`, :py:func:`~stouputils.print.error`, etc.)
message (str): Message to display with the execution time (e.g. "Execution time of Something"),
defaults to "Execution time of {func.__name__}"
perf_counter (bool): Whether to use time.perf_counter_ns or time.time_ns
defaults to True (use time.perf_counter_ns)
is_generator (bool): Whether the function is a generator or not (default: False)
When True, the decorator will yield from the function instead of returning it.
Returns:
Callable: Decorator to measure the time of the function.
Examples:
.. code-block:: python
> @measure_time(printer=info)
> def test():
> pass
> test() # [INFO HH:MM:SS] Execution time of test: 0.000ms (400ns)
"""
def decorator(
func: Callable[..., T] | Callable[..., Generator[T, None, None]]
) -> Callable[..., T] | Callable[..., Generator[T, None, None]]:
# Set the message if not specified, else use the provided one
new_msg: str = message if message else f"Execution time of {get_function_name(func)}()"
if is_generator:
@wraps(func)
def generator_wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Generator[T, None, None]:
with MeasureTime(print_func=printer, message=new_msg, perf_counter=perf_counter):
yield from func(*args, **kwargs) # type: ignore
set_wrapper_name(generator_wrapper, get_wrapper_name("stouputils.decorators.measure_time", func))
return generator_wrapper
else:
@wraps(func)
def regular_wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> T:
with MeasureTime(print_func=printer, message=new_msg, perf_counter=perf_counter):
return func(*args, **kwargs) # type: ignore
set_wrapper_name(regular_wrapper, get_wrapper_name("stouputils.decorators.measure_time", func))
return regular_wrapper
# Handle both @measure_time and @measure_time(printer=..., message=..., perf_counter=..., is_generator=...)
if func is None:
return decorator # type: ignore
return decorator(func)
