From db86c326f10c3a23817f2b80bf905e26394c0296 Mon Sep 17 00:00:00 2001
From: Marco Ribeiro <speaktomarco@gmail.com>
Date: Fri, 30 Sep 2022 15:41:21 -0300
Subject: [PATCH] Use typing.TypeVar on decorators' type hints

* Fixes issue #1654
---
 elasticapm/contrib/asyncio/traces.py | 6 ++++--
 elasticapm/contrib/serverless/aws.py | 6 ++++--
 elasticapm/traces.py                 | 5 +++--
 elasticapm/utils/compat.py           | 5 ++++-
 elasticapm/utils/deprecation.py      | 5 ++++-
 5 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/elasticapm/contrib/asyncio/traces.py b/elasticapm/contrib/asyncio/traces.py
index 1ab07aa97..34b18eabd 100644
--- a/elasticapm/contrib/asyncio/traces.py
+++ b/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)
diff --git a/elasticapm/contrib/serverless/aws.py b/elasticapm/contrib/serverless/aws.py
index 4bd2ae2f1..67d6f7211 100644
--- a/elasticapm/contrib/serverless/aws.py
+++ b/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)
diff --git a/elasticapm/traces.py b/elasticapm/traces.py
index 1a7ade1d5..b20868df6 100644
--- a/elasticapm/traces.py
+++ b/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)
diff --git a/elasticapm/utils/compat.py b/elasticapm/utils/compat.py
index f29326c6d..b0ee9876a 100644
--- a/elasticapm/utils/compat.py
+++ b/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)
diff --git a/elasticapm/utils/deprecation.py b/elasticapm/utils/deprecation.py
index 3849bd394..5c53e265f 100644
--- a/elasticapm/utils/deprecation.py
+++ b/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__)