Migrate typing for integrations - part 4

Author: sl0thentr0py

sentry_sdk/integrations/sanic.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import sys
 import weakref
 from inspect import isawaitable
@@ -59,8 +60,9 @@ class SanicIntegration(Integration):
     origin = f"auto.http.{identifier}"
     version = None
 
-    def __init__(self, unsampled_statuses=frozenset({404})):
-        # type: (Optional[Container[int]]) -> None
+    def __init__(
+        self, unsampled_statuses: Optional[Container[int]] = frozenset({404})
+    ) -> None:
         """
         The unsampled_statuses parameter can be used to specify for which HTTP statuses the
         transactions should not be sent to Sentry. By default, transactions are sent for all
@@ -70,8 +72,7 @@ def __init__(self, unsampled_statuses=frozenset({404})):
         self._unsampled_statuses = unsampled_statuses or set()
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         SanicIntegration.version = parse_version(SANIC_VERSION)
         _check_minimum_version(SanicIntegration, SanicIntegration.version)
 
@@ -103,56 +104,45 @@ def setup_once():
 
 
 class SanicRequestExtractor(RequestExtractor):
-    def content_length(self):
-        # type: () -> int
+    def content_length(self) -> int:
         if self.request.body is None:
             return 0
         return len(self.request.body)
 
-    def cookies(self):
-        # type: () -> Dict[str, str]
+    def cookies(self) -> Dict[str, str]:
         return dict(self.request.cookies)
 
-    def raw_data(self):
-        # type: () -> bytes
+    def raw_data(self) -> bytes:
         return self.request.body
 
-    def form(self):
-        # type: () -> RequestParameters
+    def form(self) -> RequestParameters:
         return self.request.form
 
-    def is_json(self):
-        # type: () -> bool
+    def is_json(self) -> bool:
         raise NotImplementedError()
 
-    def json(self):
-        # type: () -> Optional[Any]
+    def json(self) -> Optional[Any]:
         return self.request.json
 
-    def files(self):
-        # type: () -> RequestParameters
+    def files(self) -> RequestParameters:
         return self.request.files
 
-    def size_of_file(self, file):
-        # type: (Any) -> int
+    def size_of_file(self, file: Any) -> int:
         return len(file.body or ())
 
 
-def _setup_sanic():
-    # type: () -> None
+def _setup_sanic() -> None:
     Sanic._startup = _startup
     ErrorHandler.lookup = _sentry_error_handler_lookup
 
 
-def _setup_legacy_sanic():
-    # type: () -> None
+def _setup_legacy_sanic() -> None:
     Sanic.handle_request = _legacy_handle_request
     Router.get = _legacy_router_get
     ErrorHandler.lookup = _sentry_error_handler_lookup
 
 
-async def _startup(self):
-    # type: (Sanic) -> None
+async def _startup(self: Sanic) -> None:
     # This happens about as early in the lifecycle as possible, just after the
     # Request object is created. The body has not yet been consumed.
     self.signal("http.lifecycle.request")(_context_enter)
@@ -171,8 +161,7 @@ async def _startup(self):
     await old_startup(self)
 
 
-async def _context_enter(request):
-    # type: (Request) -> None
+async def _context_enter(request: Request) -> None:
     request.ctx._sentry_do_integration = (
         sentry_sdk.get_client().get_integration(SanicIntegration) is not None
     )
@@ -203,8 +192,9 @@ async def _context_enter(request):
     ).__enter__()
 
 
-async def _context_exit(request, response=None):
-    # type: (Request, Optional[BaseHTTPResponse]) -> None
+async def _context_exit(
+    request: Request, response: Optional[BaseHTTPResponse] = None
+) -> None:
     with capture_internal_exceptions():
         if not request.ctx._sentry_do_integration:
             return
@@ -233,17 +223,17 @@ async def _context_exit(request, response=None):
         request.ctx._sentry_scope_manager.__exit__(None, None, None)
 
 
-async def _set_transaction(request, route, **_):
-    # type: (Request, Route, **Any) -> None
+async def _set_transaction(request: Request, route: Route, **_: Any) -> None:
     if request.ctx._sentry_do_integration:
         with capture_internal_exceptions():
             scope = sentry_sdk.get_current_scope()
             route_name = route.name.replace(request.app.name, "").strip(".")
             scope.set_transaction_name(route_name, source=TransactionSource.COMPONENT)
 
 
-def _sentry_error_handler_lookup(self, exception, *args, **kwargs):
-    # type: (Any, Exception, *Any, **Any) -> Optional[object]
+def _sentry_error_handler_lookup(
+    self: Any, exception: Exception, *args: Any, **kwargs: Any
+) -> Optional[object]:
     _capture_exception(exception)
     old_error_handler = old_error_handler_lookup(self, exception, *args, **kwargs)
 
@@ -253,8 +243,9 @@ def _sentry_error_handler_lookup(self, exception, *args, **kwargs):
     if sentry_sdk.get_client().get_integration(SanicIntegration) is None:
         return old_error_handler
 
-    async def sentry_wrapped_error_handler(request, exception):
-        # type: (Request, Exception) -> Any
+    async def sentry_wrapped_error_handler(
+        request: Request, exception: Exception
+    ) -> Any:
         try:
             response = old_error_handler(request, exception)
             if isawaitable(response):
@@ -276,8 +267,9 @@ async def sentry_wrapped_error_handler(request, exception):
     return sentry_wrapped_error_handler
 
 
-async def _legacy_handle_request(self, request, *args, **kwargs):
-    # type: (Any, Request, *Any, **Any) -> Any
+async def _legacy_handle_request(
+    self: Any, request: Request, *args: Any, **kwargs: Any
+) -> Any:
     if sentry_sdk.get_client().get_integration(SanicIntegration) is None:
         return await old_handle_request(self, request, *args, **kwargs)
 
@@ -294,8 +286,7 @@ async def _legacy_handle_request(self, request, *args, **kwargs):
         return response
 
 
-def _legacy_router_get(self, *args):
-    # type: (Any, Union[Any, Request]) -> Any
+def _legacy_router_get(self: Any, *args: Union[Any, Request]) -> Any:
     rv = old_router_get(self, *args)
     if sentry_sdk.get_client().get_integration(SanicIntegration) is not None:
         with capture_internal_exceptions():
@@ -325,8 +316,7 @@ def _legacy_router_get(self, *args):
 
 
 @ensure_integration_enabled(SanicIntegration)
-def _capture_exception(exception):
-    # type: (Union[ExcInfo, BaseException]) -> None
+def _capture_exception(exception: Union[ExcInfo, BaseException]) -> None:
     with capture_internal_exceptions():
         event, hint = event_from_exception(
             exception,
@@ -340,10 +330,8 @@ def _capture_exception(exception):
         sentry_sdk.capture_event(event, hint=hint)
 
 
-def _make_request_processor(weak_request):
-    # type: (Callable[[], Request]) -> EventProcessor
-    def sanic_processor(event, hint):
-        # type: (Event, Optional[Hint]) -> Optional[Event]
+def _make_request_processor(weak_request: Callable[[], Request]) -> EventProcessor:
+    def sanic_processor(event: Event, hint: Optional[Hint]) -> Optional[Event]:
 
         try:
             if hint and issubclass(hint["exc_info"][0], SanicException):

sentry_sdk/integrations/serverless.py

@@ -1,47 +1,39 @@
+from __future__ import annotations
 import sys
 from functools import wraps
 
 import sentry_sdk
 from sentry_sdk.utils import event_from_exception, reraise
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
 if TYPE_CHECKING:
     from typing import Any
     from typing import Callable
     from typing import TypeVar
     from typing import Union
     from typing import Optional
-    from typing import overload
 
     F = TypeVar("F", bound=Callable[..., Any])
 
-else:
 
-    def overload(x):
-        # type: (F) -> F
-        return x
-
-
-@overload
-def serverless_function(f, flush=True):
-    # type: (F, bool) -> F
-    pass
+if TYPE_CHECKING:
 
+    @overload
+    def serverless_function(f: F, flush: bool = True) -> F:
+        pass
 
-@overload
-def serverless_function(f=None, flush=True):  # noqa: F811
-    # type: (None, bool) -> Callable[[F], F]
-    pass
+    @overload
+    def serverless_function(f: None = None, flush: bool = True) -> Callable[[F], F]:
+        pass
 
 
-def serverless_function(f=None, flush=True):  # noqa
-    # type: (Optional[F], bool) -> Union[F, Callable[[F], F]]
-    def wrapper(f):
-        # type: (F) -> F
+def serverless_function(
+    f: Optional[F] = None, flush: bool = True
+) -> Union[F, Callable[[F], F]]:
+    def wrapper(f: F) -> F:
         @wraps(f)
-        def inner(*args, **kwargs):
-            # type: (*Any, **Any) -> Any
+        def inner(*args: Any, **kwargs: Any) -> Any:
             with sentry_sdk.isolation_scope() as scope:
                 scope.clear_breadcrumbs()
 
@@ -61,8 +53,7 @@ def inner(*args, **kwargs):
         return wrapper(f)
 
 
-def _capture_and_reraise():
-    # type: () -> None
+def _capture_and_reraise() -> None:
     exc_info = sys.exc_info()
     client = sentry_sdk.get_client()
     if client.is_active():

sentry_sdk/integrations/socket.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import socket
 
 import sentry_sdk
@@ -17,17 +18,17 @@ class SocketIntegration(Integration):
     origin = f"auto.socket.{identifier}"
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         """
         patches two of the most used functions of socket: create_connection and getaddrinfo(dns resolver)
         """
         _patch_create_connection()
         _patch_getaddrinfo()
 
 
-def _get_span_description(host, port):
-    # type: (Union[bytes, str, None], Union[bytes, str, int, None]) -> str
+def _get_span_description(
+    host: Union[bytes, str, None], port: Union[bytes, str, int, None]
+) -> str:
 
     try:
         host = host.decode()  # type: ignore
@@ -43,16 +44,14 @@ def _get_span_description(host, port):
     return description
 
 
-def _patch_create_connection():
-    # type: () -> None
+def _patch_create_connection() -> None:
     real_create_connection = socket.create_connection
 
     def create_connection(
-        address,
-        timeout=socket._GLOBAL_DEFAULT_TIMEOUT,  # type: ignore
-        source_address=None,
-    ):
-        # type: (Tuple[Optional[str], int], Optional[float], Optional[Tuple[Union[bytearray, bytes, str], int]])-> socket.socket
+        address: Tuple[Optional[str], int],
+        timeout: Optional[float] = socket._GLOBAL_DEFAULT_TIMEOUT,  # type: ignore
+        source_address: Optional[Tuple[Union[bytearray, bytes, str], int]] = None,
+    ) -> socket.socket:
         integration = sentry_sdk.get_client().get_integration(SocketIntegration)
         if integration is None:
             return real_create_connection(address, timeout, source_address)
@@ -76,12 +75,25 @@ def create_connection(
     socket.create_connection = create_connection  # type: ignore
 
 
-def _patch_getaddrinfo():
-    # type: () -> None
+def _patch_getaddrinfo() -> None:
     real_getaddrinfo = socket.getaddrinfo
 
-    def getaddrinfo(host, port, family=0, type=0, proto=0, flags=0):
-        # type: (Union[bytes, str, None], Union[bytes, str, int, None], int, int, int, int) -> List[Tuple[AddressFamily, SocketKind, int, str, Union[Tuple[str, int], Tuple[str, int, int, int], Tuple[int, bytes]]]]
+    def getaddrinfo(
+        host: Union[bytes, str, None],
+        port: Union[bytes, str, int, None],
+        family: int = 0,
+        type: int = 0,
+        proto: int = 0,
+        flags: int = 0,
+    ) -> List[
+        Tuple[
+            AddressFamily,
+            SocketKind,
+            int,
+            str,
+            Union[Tuple[str, int], Tuple[str, int, int, int], Tuple[int, bytes]],
+        ]
+    ]:
         integration = sentry_sdk.get_client().get_integration(SocketIntegration)
         if integration is None:
             return real_getaddrinfo(host, port, family, type, proto, flags)