Migrate typing for serializer and misc files (#4522)

Author: sl0thentr0py

sentry_sdk/ai/monitoring.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import inspect
 from functools import wraps
 
@@ -15,22 +16,17 @@
 _ai_pipeline_name = ContextVar("ai_pipeline_name", default=None)
 
 
-def set_ai_pipeline_name(name):
-    # type: (Optional[str]) -> None
+def set_ai_pipeline_name(name: Optional[str]) -> None:
     _ai_pipeline_name.set(name)
 
 
-def get_ai_pipeline_name():
-    # type: () -> Optional[str]
+def get_ai_pipeline_name() -> Optional[str]:
     return _ai_pipeline_name.get()
 
 
-def ai_track(description, **span_kwargs):
-    # type: (str, Any) -> Callable[..., Any]
-    def decorator(f):
-        # type: (Callable[..., Any]) -> Callable[..., Any]
-        def sync_wrapped(*args, **kwargs):
-            # type: (Any, Any) -> Any
+def ai_track(description: str, **span_kwargs: Any) -> Callable[..., Any]:
+    def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+        def sync_wrapped(*args: Any, **kwargs: Any) -> Any:
             curr_pipeline = _ai_pipeline_name.get()
             op = span_kwargs.get("op", "ai.run" if curr_pipeline else "ai.pipeline")
 
@@ -60,8 +56,7 @@ def sync_wrapped(*args, **kwargs):
                         _ai_pipeline_name.set(None)
                     return res
 
-        async def async_wrapped(*args, **kwargs):
-            # type: (Any, Any) -> Any
+        async def async_wrapped(*args: Any, **kwargs: Any) -> Any:
             curr_pipeline = _ai_pipeline_name.get()
             op = span_kwargs.get("op", "ai.run" if curr_pipeline else "ai.pipeline")
 
@@ -100,9 +95,11 @@ async def async_wrapped(*args, **kwargs):
 
 
 def record_token_usage(
-    span, prompt_tokens=None, completion_tokens=None, total_tokens=None
-):
-    # type: (Span, Optional[int], Optional[int], Optional[int]) -> None
+    span: Span,
+    prompt_tokens: Optional[int] = None,
+    completion_tokens: Optional[int] = None,
+    total_tokens: Optional[int] = None,
+) -> None:
     ai_pipeline_name = get_ai_pipeline_name()
     if ai_pipeline_name:
         span.set_attribute(SPANDATA.AI_PIPELINE_NAME, ai_pipeline_name)

sentry_sdk/ai/utils.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -7,8 +8,7 @@
 from sentry_sdk.utils import logger
 
 
-def _normalize_data(data):
-    # type: (Any) -> Any
+def _normalize_data(data: Any) -> Any:
 
     # convert pydantic data (e.g. OpenAI v1+) to json compatible format
     if hasattr(data, "model_dump"):
@@ -26,7 +26,6 @@ def _normalize_data(data):
     return data
 
 
-def set_data_normalized(span, key, value):
-    # type: (Span, str, Any) -> None
+def set_data_normalized(span: Span, key: str, value: Any) -> None:
     normalized = _normalize_data(value)
     span.set_attribute(key, normalized)

sentry_sdk/crons/api.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import uuid
 
 import sentry_sdk
@@ -10,25 +11,24 @@
 
 
 def _create_check_in_event(
-    monitor_slug=None,  # type: Optional[str]
-    check_in_id=None,  # type: Optional[str]
-    status=None,  # type: Optional[str]
-    duration_s=None,  # type: Optional[float]
-    monitor_config=None,  # type: Optional[MonitorConfig]
-):
-    # type: (...) -> Event
+    monitor_slug: Optional[str] = None,
+    check_in_id: Optional[str] = None,
+    status: Optional[str] = None,
+    duration_s: Optional[float] = None,
+    monitor_config: Optional[MonitorConfig] = None,
+) -> Event:
     options = sentry_sdk.get_client().options
-    check_in_id = check_in_id or uuid.uuid4().hex  # type: str
+    check_in_id = check_in_id or uuid.uuid4().hex
 
-    check_in = {
+    check_in: Event = {
         "type": "check_in",
         "monitor_slug": monitor_slug,
         "check_in_id": check_in_id,
         "status": status,
         "duration": duration_s,
         "environment": options.get("environment", None),
         "release": options.get("release", None),
-    }  # type: Event
+    }
 
     if monitor_config:
         check_in["monitor_config"] = monitor_config
@@ -37,13 +37,12 @@ def _create_check_in_event(
 
 
 def capture_checkin(
-    monitor_slug=None,  # type: Optional[str]
-    check_in_id=None,  # type: Optional[str]
-    status=None,  # type: Optional[str]
-    duration=None,  # type: Optional[float]
-    monitor_config=None,  # type: Optional[MonitorConfig]
-):
-    # type: (...) -> str
+    monitor_slug: Optional[str] = None,
+    check_in_id: Optional[str] = None,
+    status: Optional[str] = None,
+    duration: Optional[float] = None,
+    monitor_config: Optional[MonitorConfig] = None,
+) -> str:
     check_in_event = _create_check_in_event(
         monitor_slug=monitor_slug,
         check_in_id=check_in_id,

sentry_sdk/crons/decorator.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from functools import wraps
 from inspect import iscoroutinefunction
 
@@ -16,8 +17,6 @@
         ParamSpec,
         Type,
         TypeVar,
-        Union,
-        cast,
         overload,
     )
     from sentry_sdk._types import MonitorConfig
@@ -55,22 +54,28 @@ def test(arg):
     ```
     """
 
-    def __init__(self, monitor_slug=None, monitor_config=None):
-        # type: (Optional[str], Optional[MonitorConfig]) -> None
+    def __init__(
+        self,
+        monitor_slug: Optional[str] = None,
+        monitor_config: Optional[MonitorConfig] = None,
+    ) -> None:
         self.monitor_slug = monitor_slug
         self.monitor_config = monitor_config
 
-    def __enter__(self):
-        # type: () -> None
+    def __enter__(self) -> None:
         self.start_timestamp = now()
         self.check_in_id = capture_checkin(
             monitor_slug=self.monitor_slug,
             status=MonitorStatus.IN_PROGRESS,
             monitor_config=self.monitor_config,
         )
 
-    def __exit__(self, exc_type, exc_value, traceback):
-        # type: (Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]) -> None
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
         duration_s = now() - self.start_timestamp
 
         if exc_type is None and exc_value is None and traceback is None:
@@ -89,46 +94,39 @@ def __exit__(self, exc_type, exc_value, traceback):
     if TYPE_CHECKING:
 
         @overload
-        def __call__(self, fn):
-            # type: (Callable[P, Awaitable[Any]]) -> Callable[P, Awaitable[Any]]
+        def __call__(
+            self, fn: Callable[P, Awaitable[Any]]
+        ) -> Callable[P, Awaitable[Any]]:
             # Unfortunately, mypy does not give us any reliable way to type check the
             # return value of an Awaitable (i.e. async function) for this overload,
             # since calling iscouroutinefunction narrows the type to Callable[P, Awaitable[Any]].
             ...
 
         @overload
-        def __call__(self, fn):
-            # type: (Callable[P, R]) -> Callable[P, R]
-            ...
+        def __call__(self, fn: Callable[P, R]) -> Callable[P, R]: ...
 
     def __call__(
         self,
-        fn,  # type: Union[Callable[P, R], Callable[P, Awaitable[Any]]]
-    ):
-        # type: (...) -> Union[Callable[P, R], Callable[P, Awaitable[Any]]]
+        fn: Callable[..., Any],
+    ) -> Callable[..., Any]:
         if iscoroutinefunction(fn):
             return self._async_wrapper(fn)
-
         else:
-            if TYPE_CHECKING:
-                fn = cast("Callable[P, R]", fn)
             return self._sync_wrapper(fn)
 
-    def _async_wrapper(self, fn):
-        # type: (Callable[P, Awaitable[Any]]) -> Callable[P, Awaitable[Any]]
+    def _async_wrapper(
+        self, fn: Callable[P, Awaitable[Any]]
+    ) -> Callable[P, Awaitable[Any]]:
         @wraps(fn)
-        async def inner(*args: "P.args", **kwargs: "P.kwargs"):
-            # type: (...) -> R
+        async def inner(*args: P.args, **kwargs: P.kwargs) -> R:
             with self:
                 return await fn(*args, **kwargs)
 
         return inner
 
-    def _sync_wrapper(self, fn):
-        # type: (Callable[P, R]) -> Callable[P, R]
+    def _sync_wrapper(self, fn: Callable[P, R]) -> Callable[P, R]:
         @wraps(fn)
-        def inner(*args: "P.args", **kwargs: "P.kwargs"):
-            # type: (...) -> R
+        def inner(*args: P.args, **kwargs: P.kwargs) -> R:
             with self:
                 return fn(*args, **kwargs)