From e44906e831173849dad10e701e9faf8300bb9c4e Mon Sep 17 00:00:00 2001 From: Thomas Kowalski Date: Wed, 3 Dec 2025 23:20:15 +0100 Subject: [PATCH] chore(profiling): improve typing in threading.py --- .../datadog/profiling/stack_v2/__init__.pyi | 2 +- ddtrace/internal/datadog/profiling/util.py | 10 +++---- ddtrace/profiling/collector/threading.py | 26 ++++++++++++------- 3 files changed, 22 insertions(+), 16 deletions(-) diff --git a/ddtrace/internal/datadog/profiling/stack_v2/__init__.pyi b/ddtrace/internal/datadog/profiling/stack_v2/__init__.pyi index a051cb22c3d..fb58d23dfe2 100644 --- a/ddtrace/internal/datadog/profiling/stack_v2/__init__.pyi +++ b/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: ... diff --git a/ddtrace/internal/datadog/profiling/util.py b/ddtrace/internal/datadog/profiling/util.py index c93870679a2..8a7cb609318 100644 --- a/ddtrace/internal/datadog/profiling/util.py +++ b/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 diff --git a/ddtrace/profiling/collector/threading.py b/ddtrace/profiling/collector/threading.py index dc1f1404546..9b420339d2a 100644 --- a/ddtrace/profiling/collector/threading.py +++ b/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]