Skip to content

Commit 7ab639d

Browse files
chore(profiling): improve typing in threading.py (#15504)
## Description https://datadoghq.atlassian.net/browse/PROF-13197 This PR improves typing in the `threading.py` module.
1 parent 377d52c commit 7ab639d

File tree

3 files changed

+22
-16
lines changed

3 files changed

+22
-16
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ def link_span(span: Optional[Union[context.Context, ddspan.Span]]) -> None: ...
2020

2121
# Thread management
2222
def register_thread(python_thread_id: int, native_id: int, name: str) -> None: ...
23-
def unregister_thread(name: str) -> None: ...
23+
def unregister_thread(python_thread_id: int) -> None: ...
2424

2525
# Asyncio support
2626
def track_asyncio_loop(thread_id: int, loop: Optional[asyncio.AbstractEventLoop]) -> None: ...
Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from sys import version_info
2-
from typing import Any # noqa:F401
2+
from typing import Any
33

44
from ddtrace.internal.logger import get_logger
55

@@ -8,8 +8,7 @@
88

99

1010
# 3.11 and above
11-
def _sanitize_string_check(value):
12-
# type: (Any) -> str
11+
def _sanitize_string_check(value: Any) -> str:
1312
if isinstance(value, str):
1413
return value
1514
elif value is None:
@@ -22,10 +21,9 @@ def _sanitize_string_check(value):
2221

2322

2423
# 3.10 and below (the noop version)
25-
def _sanitize_string_identity(value):
26-
# type: (Any) -> str
24+
def _sanitize_string_identity(value: Any) -> str:
2725
return value or ""
2826

2927

3028
# Assign based on version
31-
sanitize_string = _sanitize_string_check if version_info[:2] > (3, 10) else _sanitize_string_identity
29+
sanitize_string: object = _sanitize_string_check if version_info[:2] > (3, 10) else _sanitize_string_identity

ddtrace/profiling/collector/threading.py

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -60,16 +60,24 @@ class ThreadingBoundedSemaphoreCollector(_lock.LockCollector):
6060
# Also patch threading.Thread so echion can track thread lifetimes
6161
def init_stack_v2() -> None:
6262
if config.stack.enabled and stack_v2.is_available:
63-
_thread_set_native_id = ddtrace_threading.Thread._set_native_id # type: ignore[attr-defined]
64-
_thread_bootstrap_inner = ddtrace_threading.Thread._bootstrap_inner # type: ignore[attr-defined]
65-
66-
def thread_set_native_id(self, *args, **kwargs):
67-
_thread_set_native_id(self, *args, **kwargs)
68-
stack_v2.register_thread(self.ident, self.native_id, self.name)
69-
70-
def thread_bootstrap_inner(self, *args, **kwargs):
63+
_thread_set_native_id = typing.cast(
64+
typing.Callable[[threading.Thread], None],
65+
ddtrace_threading.Thread._set_native_id, # type: ignore[attr-defined]
66+
)
67+
_thread_bootstrap_inner = typing.cast(
68+
typing.Callable[[threading.Thread], None],
69+
ddtrace_threading.Thread._bootstrap_inner, # type: ignore[attr-defined]
70+
)
71+
72+
def thread_set_native_id(self: threading.Thread) -> None:
73+
_thread_set_native_id(self)
74+
if self.ident is not None and self.native_id is not None:
75+
stack_v2.register_thread(self.ident, self.native_id, self.name)
76+
77+
def thread_bootstrap_inner(self: threading.Thread, *args: typing.Any, **kwargs: typing.Any) -> None:
7178
_thread_bootstrap_inner(self, *args, **kwargs)
72-
stack_v2.unregister_thread(self.ident)
79+
if self.ident is not None:
80+
stack_v2.unregister_thread(self.ident)
7381

7482
ddtrace_threading.Thread._set_native_id = thread_set_native_id # type: ignore[attr-defined]
7583
ddtrace_threading.Thread._bootstrap_inner = thread_bootstrap_inner # type: ignore[attr-defined]

0 commit comments

Comments
 (0)