4
4
import functools
5
5
import logging
6
6
from datetime import datetime
7
- from typing import Any , Callable , Optional , Union
7
+ from typing import Callable , Optional , TypeVar , Union , overload
8
8
9
9
from samtranslator .metrics .metrics import DummyMetricsPublisher , Metrics
10
10
from samtranslator .model import Resource
11
11
12
12
LOG = logging .getLogger (__name__ )
13
13
14
+ _RT = TypeVar ("_RT" ) # return value
15
+
14
16
15
17
class MetricsMethodWrapperSingleton :
16
18
"""
@@ -76,9 +78,21 @@ def _send_cw_metric(prefix, name, execution_time_ms, func, args): # type: ignor
76
78
LOG .warning ("Failed to add metrics" , exc_info = e )
77
79
78
80
81
+ @overload
82
+ def cw_timer (
83
+ * , name : Optional [str ] = None , prefix : Optional [str ] = None
84
+ ) -> Callable [[Callable [..., _RT ]], Callable [..., _RT ]]:
85
+ ...
86
+
87
+
88
+ @overload
89
+ def cw_timer (_func : Callable [..., _RT ], name : Optional [str ] = None , prefix : Optional [str ] = None ) -> Callable [..., _RT ]:
90
+ ...
91
+
92
+
79
93
def cw_timer (
80
- _func : Optional [Callable [..., Any ]] = None , name : Optional [str ] = None , prefix : Optional [str ] = None
81
- ) -> Union [Callable [..., Any ], Callable [[Callable [..., Any ]], Callable [..., Any ]]]:
94
+ _func : Optional [Callable [..., _RT ]] = None , name : Optional [str ] = None , prefix : Optional [str ] = None
95
+ ) -> Union [Callable [..., _RT ], Callable [[Callable [..., _RT ]], Callable [..., _RT ]]]:
82
96
"""
83
97
A method decorator, that will calculate execution time of the decorated method, and store this information as a
84
98
metric in CloudWatch by calling the metrics singleton instance.
@@ -91,9 +105,9 @@ def cw_timer(
91
105
If prefix is defined, it will be added in the beginning of what is been generated above
92
106
"""
93
107
94
- def cw_timer_decorator (func : Callable [..., Any ]) -> Callable [..., Any ]:
108
+ def cw_timer_decorator (func : Callable [..., _RT ]) -> Callable [..., _RT ]:
95
109
@functools .wraps (func )
96
- def wrapper_cw_timer (* args , ** kwargs ): # type: ignore[no-untyped-def]
110
+ def wrapper_cw_timer (* args , ** kwargs ) -> _RT : # type: ignore[no-untyped-def]
97
111
start_time = datetime .now ()
98
112
99
113
exec_result = func (* args , ** kwargs )
0 commit comments