Use typing.TypeVar on decorators' type hints (#1655)

* Fixes issue #1654

Co-authored-by: Colton Myers <[email protected]>

Author: marcoffee

elasticapm/contrib/asyncio/traces.py

@@ -30,15 +30,17 @@
 
 import functools
 from types import TracebackType
-from typing import Optional, Type
+from typing import Optional, Type, TypeVar
 
 from elasticapm.conf.constants import LABEL_RE
 from elasticapm.traces import SpanType, capture_span, execution_context
 from elasticapm.utils import get_name_from_func
 
+_AnnotatedFunctionT = TypeVar("_AnnotatedFunctionT")
+
 
 class async_capture_span(capture_span):
-    def __call__(self, func):
+    def __call__(self, func: _AnnotatedFunctionT) -> _AnnotatedFunctionT:
         self.name = self.name or get_name_from_func(func)
 
         @functools.wraps(func)

elasticapm/contrib/serverless/aws.py

@@ -35,7 +35,7 @@
 import os
 import platform
 import time
-from typing import Optional
+from typing import Optional, TypeVar
 from urllib.parse import urlencode
 
 import elasticapm
@@ -51,6 +51,8 @@
 
 COLD_START = True
 
+_AnnotatedFunctionT = TypeVar("_AnnotatedFunctionT")
+
 
 class capture_serverless(object):
     """
@@ -89,7 +91,7 @@ def __init__(self, name: Optional[str] = None, elasticapm_client: Optional[Clien
 
         self.client_kwargs = kwargs
 
-    def __call__(self, func):
+    def __call__(self, func: _AnnotatedFunctionT) -> _AnnotatedFunctionT:
         self.name = self.name or get_name_from_func(func)
 
         @functools.wraps(func)

elasticapm/traces.py

@@ -39,7 +39,7 @@
 from collections import defaultdict
 from datetime import timedelta
 from types import TracebackType
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, TypeVar, Union
 
 import elasticapm
 from elasticapm.conf import constants
@@ -62,6 +62,7 @@
 execution_context = init_execution_context()
 
 SpanType = Union["Span", "DroppedSpan"]
+_AnnotatedFunctionT = TypeVar("_AnnotatedFunctionT")
 
 
 class ChildDuration(object):
@@ -1056,7 +1057,7 @@ def __init__(
         self.sync = sync
         self.links = links
 
-    def __call__(self, func: Callable) -> Callable:
+    def __call__(self, func: _AnnotatedFunctionT) -> _AnnotatedFunctionT:
         self.name = self.name or get_name_from_func(func)
 
         @functools.wraps(func)

elasticapm/utils/compat.py

@@ -33,9 +33,12 @@
 import atexit
 import functools
 import platform
+from typing import TypeVar
 
+_AnnotatedFunctionT = TypeVar("_AnnotatedFunctionT")
 
-def noop_decorator(func):
+
+def noop_decorator(func: _AnnotatedFunctionT) -> _AnnotatedFunctionT:
     @functools.wraps(func)
     def wrapped(*args, **kwargs):
         return func(*args, **kwargs)

elasticapm/utils/deprecation.py

@@ -30,6 +30,9 @@
 
 import functools
 import warnings
+from typing import TypeVar
+
+_AnnotatedFunctionT = TypeVar("_AnnotatedFunctionT")
 
 # https://wiki.python.org/moin/PythonDecoratorLibrary#Smart_deprecation_warnings_.28with_valid_filenames.2C_line_numbers.2C_etc..29
 
@@ -39,7 +42,7 @@ def deprecated(alternative=None):
     as deprecated. It will result in a warning being emitted
     when the function is used."""
 
-    def real_decorator(func):
+    def real_decorator(func: _AnnotatedFunctionT) -> _AnnotatedFunctionT:
         @functools.wraps(func)
         def new_func(*args, **kwargs):
             msg = "Call to deprecated function {0}.".format(func.__name__)