Migrate typing for integrations - part 2 (#4530)

Author: sl0thentr0py

sentry_sdk/integrations/chalice.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import sys
 from functools import wraps
 
@@ -32,8 +33,7 @@
 
 
 class EventSourceHandler(ChaliceEventSourceHandler):  # type: ignore
-    def __call__(self, event, context):
-        # type: (Any, Any) -> Any
+    def __call__(self, event: Any, context: Any) -> Any:
         client = sentry_sdk.get_client()
 
         with sentry_sdk.isolation_scope() as scope:
@@ -56,11 +56,9 @@ def __call__(self, event, context):
                 reraise(*exc_info)
 
 
-def _get_view_function_response(app, view_function, function_args):
-    # type: (Any, F, Any) -> F
+def _get_view_function_response(app: Any, view_function: F, function_args: Any) -> F:
     @wraps(view_function)
-    def wrapped_view_function(**function_args):
-        # type: (**Any) -> Any
+    def wrapped_view_function(**function_args: Any) -> Any:
         client = sentry_sdk.get_client()
         with sentry_sdk.isolation_scope() as scope:
             with capture_internal_exceptions():
@@ -99,8 +97,7 @@ class ChaliceIntegration(Integration):
     identifier = "chalice"
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
 
         version = parse_version(CHALICE_VERSION)
 
@@ -116,8 +113,9 @@ def setup_once():
                 RestAPIEventHandler._get_view_function_response
             )
 
-        def sentry_event_response(app, view_function, function_args):
-            # type: (Any, F, Dict[str, Any]) -> Any
+        def sentry_event_response(
+            app: Any, view_function: F, function_args: Dict[str, Any]
+        ) -> Any:
             wrapped_view_function = _get_view_function_response(
                 app, view_function, function_args
             )

sentry_sdk/integrations/clickhouse_driver.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
@@ -9,27 +10,13 @@
     ensure_integration_enabled,
 )
 
-from typing import TYPE_CHECKING, cast, Any, Dict, TypeVar
+from typing import TYPE_CHECKING
 
-# Hack to get new Python features working in older versions
-# without introducing a hard dependency on `typing_extensions`
-# from: https://stackoverflow.com/a/71944042/300572
 if TYPE_CHECKING:
-    from typing import ParamSpec, Callable
-else:
-    # Fake ParamSpec
-    class ParamSpec:
-        def __init__(self, _):
-            self.args = None
-            self.kwargs = None
+    from typing import ParamSpec, Callable, Any, Dict, TypeVar
 
-    # Callable[anything] will return None
-    class _Callable:
-        def __getitem__(self, _):
-            return None
-
-    # Make instances
-    Callable = _Callable()
+    P = ParamSpec("P")
+    T = TypeVar("T")
 
 
 try:
@@ -72,10 +59,6 @@ def setup_once() -> None:
         )
 
 
-P = ParamSpec("P")
-T = TypeVar("T")
-
-
 def _wrap_start(f: Callable[P, T]) -> Callable[P, T]:
     @ensure_integration_enabled(ClickhouseDriverIntegration, f)
     def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
@@ -93,8 +76,7 @@ def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
 
         connection._sentry_span = span  # type: ignore[attr-defined]
 
-        data = _get_db_data(connection)
-        data = cast("dict[str, Any]", data)
+        data: dict[str, Any] = _get_db_data(connection)
         data["db.query.text"] = query
 
         if query_id:
@@ -117,7 +99,11 @@ def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
 def _wrap_end(f: Callable[P, T]) -> Callable[P, T]:
     def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
         res = f(*args, **kwargs)
-        client = cast("clickhouse_driver.client.Client", args[0])
+
+        client = args[0]
+        if not isinstance(client, clickhouse_driver.client.Client):
+            return res
+
         connection = client.connection
 
         span = getattr(connection, "_sentry_span", None)
@@ -150,21 +136,25 @@ def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
 
 def _wrap_send_data(f: Callable[P, T]) -> Callable[P, T]:
     def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T:
-        client = cast("clickhouse_driver.client.Client", args[0])
+        client = args[0]
+        if not isinstance(client, clickhouse_driver.client.Client):
+            return f(*args, **kwargs)
+
         connection = client.connection
-        db_params_data = cast("list[Any]", args[2])
         span = getattr(connection, "_sentry_span", None)
 
         if span is not None:
             data = _get_db_data(connection)
             _set_on_span(span, data)
 
             if should_send_default_pii():
-                saved_db_data = getattr(
+                saved_db_data: dict[str, Any] = getattr(
                     connection, "_sentry_db_data", {}
-                )  # type: dict[str, Any]
-                db_params = saved_db_data.get("db.params") or []  # type: list[Any]
-                db_params.extend(db_params_data)
+                )
+                db_params: list[Any] = saved_db_data.get("db.params") or []
+                db_params_data = args[2]
+                if isinstance(db_params_data, list):
+                    db_params.extend(db_params_data)
                 saved_db_data["db.params"] = db_params
                 span.set_attribute("db.params", _serialize_span_attribute(db_params))
 

sentry_sdk/integrations/cloud_resource_context.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import json
 import urllib3
 
@@ -65,13 +66,11 @@ class CloudResourceContextIntegration(Integration):
 
     gcp_metadata = None
 
-    def __init__(self, cloud_provider=""):
-        # type: (str) -> None
+    def __init__(self, cloud_provider: str = "") -> None:
         CloudResourceContextIntegration.cloud_provider = cloud_provider
 
     @classmethod
-    def _is_aws(cls):
-        # type: () -> bool
+    def _is_aws(cls) -> bool:
         try:
             r = cls.http.request(
                 "PUT",
@@ -95,8 +94,7 @@ def _is_aws(cls):
             return False
 
     @classmethod
-    def _get_aws_context(cls):
-        # type: () -> Dict[str, str]
+    def _get_aws_context(cls) -> Dict[str, str]:
         ctx = {
             "cloud.provider": CLOUD_PROVIDER.AWS,
             "cloud.platform": CLOUD_PLATFORM.AWS_EC2,
@@ -149,8 +147,7 @@ def _get_aws_context(cls):
         return ctx
 
     @classmethod
-    def _is_gcp(cls):
-        # type: () -> bool
+    def _is_gcp(cls) -> bool:
         try:
             r = cls.http.request(
                 "GET",
@@ -174,8 +171,7 @@ def _is_gcp(cls):
             return False
 
     @classmethod
-    def _get_gcp_context(cls):
-        # type: () -> Dict[str, str]
+    def _get_gcp_context(cls) -> Dict[str, str]:
         ctx = {
             "cloud.provider": CLOUD_PROVIDER.GCP,
             "cloud.platform": CLOUD_PLATFORM.GCP_COMPUTE_ENGINE,
@@ -229,8 +225,7 @@ def _get_gcp_context(cls):
         return ctx
 
     @classmethod
-    def _get_cloud_provider(cls):
-        # type: () -> str
+    def _get_cloud_provider(cls) -> str:
         if cls._is_aws():
             return CLOUD_PROVIDER.AWS
 
@@ -240,8 +235,7 @@ def _get_cloud_provider(cls):
         return ""
 
     @classmethod
-    def _get_cloud_resource_context(cls):
-        # type: () -> Dict[str, str]
+    def _get_cloud_resource_context(cls) -> Dict[str, str]:
         cloud_provider = (
             cls.cloud_provider
             if cls.cloud_provider != ""
@@ -253,8 +247,7 @@ def _get_cloud_resource_context(cls):
         return {}
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         cloud_provider = CloudResourceContextIntegration.cloud_provider
         unsupported_cloud_provider = (
             cloud_provider != "" and cloud_provider not in context_getters.keys()

sentry_sdk/integrations/cohere.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from functools import wraps
 
 from sentry_sdk import consts
@@ -70,20 +71,17 @@ class CohereIntegration(Integration):
     identifier = "cohere"
     origin = f"auto.ai.{identifier}"
 
-    def __init__(self, include_prompts=True):
-        # type: (CohereIntegration, bool) -> None
+    def __init__(self: CohereIntegration, include_prompts: bool = True) -> None:
         self.include_prompts = include_prompts
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False)
         Client.embed = _wrap_embed(Client.embed)
         BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True)
 
 
-def _capture_exception(exc):
-    # type: (Any) -> None
+def _capture_exception(exc: Any) -> None:
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,
@@ -92,11 +90,11 @@ def _capture_exception(exc):
     sentry_sdk.capture_event(event, hint=hint)
 
 
-def _wrap_chat(f, streaming):
-    # type: (Callable[..., Any], bool) -> Callable[..., Any]
+def _wrap_chat(f: Callable[..., Any], streaming: bool) -> Callable[..., Any]:
 
-    def collect_chat_response_fields(span, res, include_pii):
-        # type: (Span, NonStreamedChatResponse, bool) -> None
+    def collect_chat_response_fields(
+        span: Span, res: NonStreamedChatResponse, include_pii: bool
+    ) -> None:
         if include_pii:
             if hasattr(res, "text"):
                 set_data_normalized(
@@ -130,8 +128,7 @@ def collect_chat_response_fields(span, res, include_pii):
                 set_data_normalized(span, SPANDATA.AI_WARNINGS, res.meta.warnings)
 
     @wraps(f)
-    def new_chat(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
+    def new_chat(*args: Any, **kwargs: Any) -> Any:
         integration = sentry_sdk.get_client().get_integration(CohereIntegration)
 
         if (
@@ -185,8 +182,7 @@ def new_chat(*args, **kwargs):
             if streaming:
                 old_iterator = res
 
-                def new_iterator():
-                    # type: () -> Iterator[StreamedChatResponse]
+                def new_iterator() -> Iterator[StreamedChatResponse]:
 
                     with capture_internal_exceptions():
                         for x in old_iterator:
@@ -220,12 +216,10 @@ def new_iterator():
     return new_chat
 
 
-def _wrap_embed(f):
-    # type: (Callable[..., Any]) -> Callable[..., Any]
+def _wrap_embed(f: Callable[..., Any]) -> Callable[..., Any]:
 
     @wraps(f)
-    def new_embed(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
+    def new_embed(*args: Any, **kwargs: Any) -> Any:
         integration = sentry_sdk.get_client().get_integration(CohereIntegration)
         if integration is None:
             return f(*args, **kwargs)

sentry_sdk/integrations/dedupe.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import sentry_sdk
 from sentry_sdk.utils import ContextVar
 from sentry_sdk.integrations import Integration
@@ -14,16 +15,13 @@
 class DedupeIntegration(Integration):
     identifier = "dedupe"
 
-    def __init__(self):
-        # type: () -> None
+    def __init__(self) -> None:
         self._last_seen = ContextVar("last-seen")
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         @add_global_event_processor
-        def processor(event, hint):
-            # type: (Event, Optional[Hint]) -> Optional[Event]
+        def processor(event: Event, hint: Optional[Hint]) -> Optional[Event]:
             if hint is None:
                 return event
 
@@ -42,8 +40,7 @@ def processor(event, hint):
             return event
 
     @staticmethod
-    def reset_last_seen():
-        # type: () -> None
+    def reset_last_seen() -> None:
         integration = sentry_sdk.get_client().get_integration(DedupeIntegration)
         if integration is None:
             return