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

Author: hossein-raeisi

mypy.ini

@@ -67,3 +67,5 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-arq.*]
 ignore_missing_imports = True
+[mypy-grpc.*]
+ignore_missing_imports = True

sentry_sdk/integrations/grpc/client.py

@@ -4,10 +4,13 @@
 from sentry_sdk.integrations import DidNotEnable
 
 if MYPY:
-    pass
+    from typing import Callable, Iterator, Union
 
 try:
     import grpc
+    from grpc import ClientCallDetails, Call
+    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, Message], _UnaryOutcome], 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, Message], Iterator], ClientCallDetails, Message) -> Union[Iterator[Message], Call]
         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)
@@ -55,6 +60,7 @@ def intercept_unary_stream(self, continuation, client_call_details, request):
 
     @staticmethod
     def _update_client_call_details_metadata_from_hub(client_call_details, hub):
+        # type: (ClientCallDetails, Hub) -> ClientCallDetails
         metadata = (
             list(client_call_details.metadata) if client_call_details.metadata else []
         )

sentry_sdk/integrations/grpc/server.py

@@ -5,27 +5,31 @@
 from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_CUSTOM
 
 if MYPY:
-    from typing import Callable
+    from typing import Callable, Optional
+    from google.protobuf.message import Message
 
 try:
     import grpc
+    from grpc import ServicerContext, HandlerCallDetails, RpcMethodHandler
 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], str]]) -> 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], RpcMethodHandler], HandlerCallDetails) -> RpcMethodHandler
         handler = continuation(handler_call_details)
         if not handler or not handler.unary_unary:
             return handler
 
         def behavior(request, context):
+            # type: (Message, ServicerContext) -> Message
             hub = Hub(Hub.current)
 
             name = self._find_method_name(context)
@@ -56,4 +60,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,14 @@
 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 socket import AddressFamily, SocketKind
+    from typing import Tuple, Optional, Union, List
+
 __all__ = ["SocketIntegration"]
 
 
@@ -25,17 +30,17 @@ def _patch_create_connection():
     real_create_connection = socket.create_connection
 
     def create_connection(
-        address, timeout=socket._GLOBAL_DEFAULT_TIMEOUT, source_address=None
+            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(
                 address=address, timeout=timeout, source_address=source_address
             )
 
         with hub.start_span(
-            op=OP.SOCKET_CONNECTION, description="%s:%s" % (address[0], address[1])
+                op=OP.SOCKET_CONNECTION, description="%s:%s" % (address[0], address[1])
         ) as span:
             span.set_data("address", address)
             span.set_data("timeout", timeout)
@@ -52,17 +57,18 @@ def _patch_getaddrinfo():
     # type: () -> None
     real_getaddrinfo = socket.getaddrinfo
 
-    def getaddrinfo(host, port, *args, **kwargs):
+    def getaddrinfo(host, port, family=0, type=0, proto=0, flags=0):
+        # type: (Union[bytes, str, None], Union[str, int, None], int, int, int, int) -> List[Tuple[AddressFamily, SocketKind, int, str, Union[Tuple[str, int], Tuple[str, int, int, int]]]]
         hub = Hub.current
         if hub.get_integration(SocketIntegration) is None:
-            return real_getaddrinfo(host, port, *args, **kwargs)
+            return real_getaddrinfo(host, port, family, type, proto, flags)
 
         with hub.start_span(
-            op=OP.SOCKET_DNS, description="%s:%s" % (host, port)
+                op=OP.SOCKET_DNS, description="%s:%s" % (host, port)
         ) as span:
             span.set_data("host", host)
             span.set_data("port", port)
 
-            return real_getaddrinfo(host, port, *args, **kwargs)
+            return real_getaddrinfo(host, port, family, type, proto, flags)
 
     socket.getaddrinfo = getaddrinfo