feat(tracing) Add type hints for grpc and socket integrations

Author: hossein-raeisi

sentry_sdk/integrations/grpc/client.py

@@ -4,10 +4,13 @@
 from sentry_sdk.integrations import DidNotEnable
 
 if MYPY:
-    pass
+    from typing import Callable, Iterator
 
 try:
     import grpc
+    from grpc import ClientCallDetails
+    from grpc._interceptor import _UnaryOutcome
+    from google.protobuf.message import Message
 except ImportError:
     raise DidNotEnable("grpcio is not installed")
 
@@ -16,11 +19,12 @@ class ClientInterceptor(
     grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor
 ):
     def intercept_unary_unary(self, continuation, client_call_details, request):
+        # type: (ClientInterceptor, Callable[ClientCallDetails], ClientCallDetails, Message) -> _UnaryOutcome
         hub = Hub.current
         method = client_call_details.method
 
         with hub.start_span(
-            op=OP.GRPC_CLIENT, description="unary unary call to %s" % method
+                op=OP.GRPC_CLIENT, description="unary unary call to %s" % method
         ) as span:
             span.set_data("type", "unary unary")
             span.set_data("method", method)
@@ -35,11 +39,12 @@ def intercept_unary_unary(self, continuation, client_call_details, request):
             return response
 
     def intercept_unary_stream(self, continuation, client_call_details, request):
+        # type: (ClientInterceptor, Callable[ClientCallDetails], ClientCallDetails, Message) -> Iterator
         hub = Hub.current
         method = client_call_details.method
 
         with hub.start_span(
-            op=OP.GRPC_CLIENT, description="unary stream call to %s" % method
+                op=OP.GRPC_CLIENT, description="unary stream call to %s" % method
         ) as span:
             span.set_data("type", "unary stream")
             span.set_data("method", method)

sentry_sdk/integrations/grpc/server.py

@@ -5,22 +5,24 @@
 from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_CUSTOM
 
 if MYPY:
-    from typing import Callable
+    from typing import Callable, Optional
 
 try:
     import grpc
+    from grpc import ServicerContext, HandlerCallDetails
 except ImportError:
     raise DidNotEnable("grpcio is not installed")
 
 
 class ServerInterceptor(grpc.ServerInterceptor):
     def __init__(self, find_name=None):
-        # type: (ServerInterceptor, Callable | None) -> ServerInterceptor
+        # type: (ServerInterceptor, Optional[Callable[ServicerContext]]) -> None
         if find_name:
             self._find_method_name = find_name
         super(ServerInterceptor, self).__init__()
 
     def intercept_service(self, continuation, handler_call_details):
+        # type: (ServerInterceptor, Callable[HandlerCallDetails], HandlerCallDetails) -> RpcMethodHandler
         handler = continuation(handler_call_details)
         if not handler or not handler.unary_unary:
             return handler
@@ -56,4 +58,5 @@ def behavior(request, context):
 
     @staticmethod
     def _find_name(context):
+        # type: (ServicerContext) -> str
         return context._rpc_event.call_details.method.decode()

sentry_sdk/integrations/socket.py

@@ -1,9 +1,13 @@
 import socket
 
 from sentry_sdk import Hub
+from sentry_sdk._types import MYPY
 from sentry_sdk.consts import OP
 from sentry_sdk.integrations import Integration
 
+if MYPY:
+    from typing import Tuple, Optional, Union, Dict, List
+
 __all__ = ["SocketIntegration"]
 
 
@@ -27,7 +31,7 @@ def _patch_create_connection():
     def create_connection(
         address, timeout=socket._GLOBAL_DEFAULT_TIMEOUT, source_address=None
     ):
-        # type: (tuple[str, None | int], float | None, tuple[bytearray | bytes | str, int] | None) -> socket
+        # type: (Tuple[Optional[str], int], Optional[float], Optional[Tuple[Union[bytearray, bytes, str], int]])-> socket.socket
         hub = Hub.current
         if hub.get_integration(SocketIntegration) is None:
             return real_create_connection(
@@ -53,6 +57,7 @@ def _patch_getaddrinfo():
     real_getaddrinfo = socket.getaddrinfo
 
     def getaddrinfo(host, port, *args, **kwargs):
+        # type: (str, int, Optional[Tuple], Optional[Dict]) -> List
         hub = Hub.current
         if hub.get_integration(SocketIntegration) is None:
             return real_getaddrinfo(host, port, *args, **kwargs)