Fix typing

Author: alexander-alderman-webb

sentry_sdk/integrations/concurrent.py

@@ -1,6 +1,6 @@
 from functools import wraps
 
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, Future
 
 import sentry_sdk
 from sentry_sdk.integrations import Integration
@@ -9,8 +9,9 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any
-    from typing import Callable
+    from typing import Any, Callable, TypeVar
+
+    T = TypeVar("T", bound=Any)
 
 
 class ConcurrentIntegration(Integration):
@@ -23,24 +24,31 @@ def __init__(self, record_exceptions_on_futures=True):
     @staticmethod
     def setup_once():
         # type: () -> None
-        old_submit = ThreadPoolExecutor.submit
+        ThreadPoolExecutor.submit = _wrap_submit_call(ThreadPoolExecutor.submit)
+
+
+def _wrap_submit_call(func):
+    # type: (Any) -> Any
+    """
+    Wrap task call with a try catch to get exceptions.
+    """
 
-        @wraps(old_submit)
-        def sentry_submit(self, fn, *args, **kwargs):
-            # type: (ThreadPoolExecutor, Callable, *Any, **Any) -> Any
-            integration = sentry_sdk.get_client().get_integration(ConcurrentIntegration)
-            if integration is None:
-                return old_submit(self, fn, *args, **kwargs)
+    @wraps(func)
+    def sentry_submit(self, fn, *args, **kwargs):
+        # type: (ThreadPoolExecutor, Callable[..., T], *Any, **Any) -> Future[T]
+        integration = sentry_sdk.get_client().get_integration(ConcurrentIntegration)
+        if integration is None:
+            return func(self, fn, *args, **kwargs)
 
-            isolation_scope = sentry_sdk.get_isolation_scope().fork()
-            current_scope = sentry_sdk.get_current_scope().fork()
+        isolation_scope = sentry_sdk.get_isolation_scope().fork()
+        current_scope = sentry_sdk.get_current_scope().fork()
 
-            def wrapped_fn(*args, **kwargs):
-                # type: (*Any, **Any) -> Any
-                with use_isolation_scope(isolation_scope):
-                    with use_scope(current_scope):
-                        return fn(*args, **kwargs)
+        def wrapped_fn(*args, **kwargs):
+            # type: (*Any, **Any) -> Any
+            with use_isolation_scope(isolation_scope):
+                with use_scope(current_scope):
+                    return fn(*args, **kwargs)
 
-            return old_submit(self, wrapped_fn, *args, **kwargs)
+        return func(self, wrapped_fn, *args, **kwargs)
 
-        ThreadPoolExecutor.submit = sentry_submit
+    return sentry_submit