fix(rnd): Fix decorator function type hint (#8043)

This commit is contained in:
Zamil Majdy
2024-09-11 17:35:33 -05:00
committed by GitHub
parent c8fbce643e
commit 52c731abd6

View File

@@ -2,7 +2,7 @@ import functools
import logging
import os
import time
from typing import Callable, Tuple, TypeVar
from typing import Callable, ParamSpec, Tuple, TypeVar
from pydantic import BaseModel
@@ -24,18 +24,19 @@ def _end_measurement(
return end_wall_time - start_wall_time, end_cpu_time - start_cpu_time
P = ParamSpec("P")
T = TypeVar("T")
logger = logging.getLogger(__name__)
def time_measured(func: Callable[..., T]) -> Callable[..., Tuple[TimingInfo, T]]:
def time_measured(func: Callable[P, T]) -> Callable[P, Tuple[TimingInfo, T]]:
"""
Decorator to measure the time taken by a function to execute.
"""
@functools.wraps(func)
def wrapper(*args, **kwargs) -> Tuple[TimingInfo, T]:
def wrapper(*args: P.args, **kwargs: P.kwargs) -> Tuple[TimingInfo, T]:
start_wall_time, start_cpu_time = _start_measurement()
try:
result = func(*args, **kwargs)
@@ -49,13 +50,13 @@ def time_measured(func: Callable[..., T]) -> Callable[..., Tuple[TimingInfo, T]]
return wrapper
def error_logged(func: Callable[..., T]) -> Callable[..., T | None]:
def error_logged(func: Callable[P, T]) -> Callable[P, T | None]:
"""
Decorator to suppress and log any exceptions raised by a function.
"""
@functools.wraps(func)
def wrapper(*args, **kwargs) -> T | None:
def wrapper(*args: P.args, **kwargs: P.kwargs) -> T | None:
try:
return func(*args, **kwargs)
except Exception as e: