From 896b37d2c4497ed7dc555c75429d0d461ff56259 Mon Sep 17 00:00:00 2001 From: Matt Wozniski Date: Fri, 8 Mar 2024 17:10:21 -0500 Subject: [PATCH] Put thread name formatting in a free function Introduce a new `memray.reporters.common` module for utilities needed by multiple reporters. Replace the `AllocationRecord.pretty_thread_name` property with a `format_thread_name` function in that new module. Signed-off-by: Matt Wozniski --- src/memray/_memray.pyi | 4 ---- src/memray/_memray.pyx | 14 -------------- src/memray/reporters/common.py | 14 ++++++++++++++ src/memray/reporters/flamegraph.py | 7 ++++--- src/memray/reporters/table.py | 3 ++- src/memray/reporters/transform.py | 3 ++- src/memray/reporters/tree.py | 3 ++- tests/utils.py | 4 ---- 8 files changed, 24 insertions(+), 28 deletions(-) create mode 100644 src/memray/reporters/common.py diff --git a/src/memray/_memray.pyi b/src/memray/_memray.pyi index cc4ace2d00..90411c2a39 100644 --- a/src/memray/_memray.pyi +++ b/src/memray/_memray.pyi @@ -47,8 +47,6 @@ class AllocationRecord: def native_segment_generation(self) -> int: ... @property def thread_name(self) -> str: ... - @property - def pretty_thread_name(self) -> str: ... def hybrid_stack_trace( self, max_stacks: Optional[int] = None, @@ -94,8 +92,6 @@ class TemporalAllocationRecord: def native_segment_generation(self) -> int: ... @property def thread_name(self) -> str: ... - @property - def pretty_thread_name(self) -> str: ... def hybrid_stack_trace( self, max_stacks: Optional[int] = None, diff --git a/src/memray/_memray.pyx b/src/memray/_memray.pyx index f149431ed8..2e1b7ef77f 100644 --- a/src/memray/_memray.pyx +++ b/src/memray/_memray.pyx @@ -309,14 +309,6 @@ cdef class AllocationRecord: assert self._reader.get() != NULL, "Cannot get thread name without reader." return self._reader.get().getThreadName(self.tid) - @property - def pretty_thread_name(self): - if self.tid == -1: - return "merged thread" - name = self.thread_name - thread_id = hex(self.tid) - return f"{thread_id} ({name})" if name else f"{thread_id}" - def stack_trace(self, max_stacks=None): cache_key = ("python", max_stacks) if cache_key not in self._stack_trace_cache: @@ -450,12 +442,6 @@ cdef class TemporalAllocationRecord: assert self._reader.get() != NULL, "Cannot get thread name without reader." return self._reader.get().getThreadName(self.tid) - @property - def pretty_thread_name(self): - name = self.thread_name - thread_id = hex(self.tid) - return f"{thread_id} ({name})" if name else f"{thread_id}" - def stack_trace(self, max_stacks=None): cache_key = ("python", max_stacks) if cache_key not in self._stack_trace_cache: diff --git a/src/memray/reporters/common.py b/src/memray/reporters/common.py new file mode 100644 index 0000000000..c089d01529 --- /dev/null +++ b/src/memray/reporters/common.py @@ -0,0 +1,14 @@ +from typing import Union + +from memray._memray import AllocationRecord +from memray._memray import TemporalAllocationRecord + + +def format_thread_name( + record: Union[AllocationRecord, TemporalAllocationRecord] +) -> str: + if record.tid == -1: + return "merged thread" + name = record.thread_name + thread_id = hex(record.tid) + return f"{thread_id} ({name})" if name else f"{thread_id}" diff --git a/src/memray/reporters/flamegraph.py b/src/memray/reporters/flamegraph.py index 33dbf5a1fa..1f24b61e94 100644 --- a/src/memray/reporters/flamegraph.py +++ b/src/memray/reporters/flamegraph.py @@ -26,6 +26,7 @@ from memray import Metadata from memray._memray import Interval from memray._memray import TemporalAllocationRecord +from memray.reporters.common import format_thread_name from memray.reporters.frame_tools import StackFrame from memray.reporters.frame_tools import is_cpython_internal from memray.reporters.frame_tools import is_frame_from_import_system @@ -263,13 +264,13 @@ def _from_any_snapshot( unique_threads: Set[str] = set() for record in allocations: - unique_threads.add(record.pretty_thread_name) + unique_threads.add(format_thread_name(record)) record_data: RecordData if temporal: assert isinstance(record, TemporalAllocationRecord) record_data = { - "thread_name": record.pretty_thread_name, + "thread_name": format_thread_name(record), "intervals": record.intervals, "size": None, "n_allocations": None, @@ -277,7 +278,7 @@ def _from_any_snapshot( else: assert not isinstance(record, TemporalAllocationRecord) record_data = { - "thread_name": record.pretty_thread_name, + "thread_name": format_thread_name(record), "intervals": None, "size": record.size, "n_allocations": record.n_allocations, diff --git a/src/memray/reporters/table.py b/src/memray/reporters/table.py index f798477cb8..3b2ebe85ab 100644 --- a/src/memray/reporters/table.py +++ b/src/memray/reporters/table.py @@ -9,6 +9,7 @@ from memray import AllocatorType from memray import MemorySnapshot from memray import Metadata +from memray.reporters.common import format_thread_name from memray.reporters.templates import render_report @@ -47,7 +48,7 @@ def from_snapshot( allocator = AllocatorType(record.allocator) result.append( { - "tid": record.pretty_thread_name, + "tid": format_thread_name(record), "size": record.size, "allocator": allocator.name.lower(), "n_allocations": record.n_allocations, diff --git a/src/memray/reporters/transform.py b/src/memray/reporters/transform.py index 30a40ac845..5a8c821aec 100644 --- a/src/memray/reporters/transform.py +++ b/src/memray/reporters/transform.py @@ -11,6 +11,7 @@ from memray import AllocatorType from memray import MemorySnapshot from memray import Metadata +from memray.reporters.common import format_thread_name Location = Tuple[str, str] @@ -117,7 +118,7 @@ def render_as_csv( record.n_allocations, record.size, record.tid, - record.pretty_thread_name, + format_thread_name(record), "|".join(f"{func};{mod};{line}" for func, mod, line in stack_trace), ] ) diff --git a/src/memray/reporters/tree.py b/src/memray/reporters/tree.py index 8577fc8523..035d5d8210 100644 --- a/src/memray/reporters/tree.py +++ b/src/memray/reporters/tree.py @@ -40,6 +40,7 @@ from memray.reporters._textual_hacks import Bindings from memray.reporters._textual_hacks import redraw_footer from memray.reporters._textual_hacks import update_key_description +from memray.reporters.common import format_thread_name from memray.reporters.frame_tools import is_cpython_internal from memray.reporters.frame_tools import is_frame_from_import_system from memray.reporters.frame_tools import is_frame_interesting @@ -476,7 +477,7 @@ def from_snapshot( current_frame = current_frame.children[stack_frame] current_frame.value += size current_frame.n_allocations += record.n_allocations - current_frame.thread_id = record.pretty_thread_name + current_frame.thread_id = format_thread_name(record) if index > MAX_STACKS: break diff --git a/tests/utils.py b/tests/utils.py index c9000e7786..0f8d866339 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -66,10 +66,6 @@ class MockAllocationRecord: _hybrid_stack: Optional[List[Tuple[str, str, int]]] = None thread_name: str = "" - @property - def pretty_thread_name(self): - return str(hex(self.tid)) if self.tid != -1 else "merged thread" - @staticmethod def __get_stack_trace(stack, max_stacks): if max_stacks == 0: