chore(profiling): improve typing in threading.py

Author: KowalskiThomas

ddtrace/internal/datadog/profiling/stack_v2/__init__.pyi

@@ -20,7 +20,7 @@ def link_span(span: Optional[Union[context.Context, ddspan.Span]]) -> None: ...
 
 # Thread management
 def register_thread(python_thread_id: int, native_id: int, name: str) -> None: ...
-def unregister_thread(name: str) -> None: ...
+def unregister_thread(python_thread_id: int) -> None: ...
 
 # Asyncio support
 def track_asyncio_loop(thread_id: int, loop: Optional[asyncio.AbstractEventLoop]) -> None: ...

ddtrace/internal/datadog/profiling/util.py

@@ -1,5 +1,5 @@
 from sys import version_info
-from typing import Any  # noqa:F401
+from typing import Any
 
 from ddtrace.internal.logger import get_logger
 
@@ -8,8 +8,7 @@
 
 
 # 3.11 and above
-def _sanitize_string_check(value):
-    # type: (Any) -> str
+def _sanitize_string_check(value: Any) -> str:
     if isinstance(value, str):
         return value
     elif value is None:
@@ -22,10 +21,9 @@ def _sanitize_string_check(value):
 
 
 # 3.10 and below (the noop version)
-def _sanitize_string_identity(value):
-    # type: (Any) -> str
+def _sanitize_string_identity(value: Any) -> str:
     return value or ""
 
 
 # Assign based on version
-sanitize_string = _sanitize_string_check if version_info[:2] > (3, 10) else _sanitize_string_identity
+sanitize_string: object = _sanitize_string_check if version_info[:2] > (3, 10) else _sanitize_string_identity

ddtrace/profiling/collector/threading.py

@@ -51,16 +51,24 @@ def _set_patch_target(
 # Also patch threading.Thread so echion can track thread lifetimes
 def init_stack_v2() -> None:
     if config.stack.enabled and stack_v2.is_available:
-        _thread_set_native_id = ddtrace_threading.Thread._set_native_id  # type: ignore[attr-defined]
-        _thread_bootstrap_inner = ddtrace_threading.Thread._bootstrap_inner  # type: ignore[attr-defined]
-
-        def thread_set_native_id(self, *args, **kwargs):
-            _thread_set_native_id(self, *args, **kwargs)
-            stack_v2.register_thread(self.ident, self.native_id, self.name)
-
-        def thread_bootstrap_inner(self, *args, **kwargs):
+        _thread_set_native_id = typing.cast(
+            typing.Callable[[threading.Thread], None],
+            ddtrace_threading.Thread._set_native_id,  # type: ignore[attr-defined]
+        )
+        _thread_bootstrap_inner = typing.cast(
+            typing.Callable[[threading.Thread], None],
+            ddtrace_threading.Thread._bootstrap_inner,  # type: ignore[attr-defined]
+        )
+
+        def thread_set_native_id(self: threading.Thread) -> None:
+            _thread_set_native_id(self)
+            if self.ident is not None and self.native_id is not None:
+                stack_v2.register_thread(self.ident, self.native_id, self.name)
+
+        def thread_bootstrap_inner(self: threading.Thread, *args: typing.Any, **kwargs: typing.Any) -> None:
             _thread_bootstrap_inner(self, *args, **kwargs)
-            stack_v2.unregister_thread(self.ident)
+            if self.ident is not None:
+                stack_v2.unregister_thread(self.ident)
 
         ddtrace_threading.Thread._set_native_id = thread_set_native_id  # type: ignore[attr-defined]
         ddtrace_threading.Thread._bootstrap_inner = thread_bootstrap_inner  # type: ignore[attr-defined]