Skip to content

Commit

Permalink
ref(api): Improve sentry_sdk.trace type hints (#2633)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
szokeasaurusrex authored Jan 16, 2024
1 parent fe1f01b commit 2f05ccb
Showing 1 changed file with 21 additions and 1 deletion.
22 changes: 21 additions & 1 deletion sentry_sdk/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 2f05ccb

Please sign in to comment.