ref(api): Improve sentry_sdk.trace type hints (#2633)

Type hints for sentry_sdk.trace decorator function now indicate that the decorator returns a function with the same signature as it was called with. Previously, the type hints indicated that the decorator could return Any, which caused users to lose type hints for decorated functions.

* Improve `sentry_sdk.trace` type hints

* Add overloads for None case

* Fix typing when `trace` called with `None`

Fixes GH-2460

Author: szokeasaurusrex

sentry_sdk/tracing.py

@@ -14,13 +14,20 @@
 if TYPE_CHECKING:
     import typing
 
+    from collections.abc import Callable
     from typing import Any
     from typing import Dict
     from typing import Iterator
     from typing import List
     from typing import Optional
+    from typing import overload
+    from typing import ParamSpec
     from typing import Tuple
     from typing import Union
+    from typing import TypeVar
+
+    P = ParamSpec("P")
+    R = TypeVar("R")
 
     import sentry_sdk.profiler
     from sentry_sdk._types import Event, MeasurementUnit, SamplingContext
@@ -983,8 +990,21 @@ def _set_initial_sampling_decision(self, sampling_context):
         pass
 
 
+if TYPE_CHECKING:
+
+    @overload
+    def trace(func=None):
+        # type: (None) -> Callable[[Callable[P, R]], Callable[P, R]]
+        pass
+
+    @overload
+    def trace(func):
+        # type: (Callable[P, R]) -> Callable[P, R]
+        pass
+
+
 def trace(func=None):
-    # type: (Any) -> Any
+    # type: (Optional[Callable[P, R]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
     """
     Decorator to start a child span under the existing current transaction.
     If there is no current transaction, then nothing will be traced.