Add per-module allocation statistics to memray stats

Display the top N modules responsible for the most allocations in
the memray stats reporter. For each allocation, walk the call stack
to locate the nearest non-stdlib Python frame and extract its
top-level module name. Aggregate allocation counts and allocated
bytes by module, and include the top N results in both terminal
and JSON output.

Signed-off-by: Ivona Stojanovic <stojanovic.i@hotmail.com>

Author: ivonastojanovic

news/765.feature.rst

@@ -0,0 +1 @@
+Add per-module allocation statistics to ``memray stats``, showing the top allocating non-stdlib modules by size and by number of allocations.

src/memray/_memray.pyx

@@ -1493,6 +1493,9 @@ def compute_statistics(
     report_progress=False,
     num_largest=5,
 ):
+    from memray.reporters.module_tools import get_module_for_stack
+    from memray.reporters.module_tools import get_python_path_info
+
     cdef shared_ptr[RecordReader] reader_sp = make_shared[RecordReader](
         unique_ptr[FileSource](new FileSource(file_name))
     )
@@ -1512,15 +1515,27 @@ def compute_statistics(
         total=total,
         report_progress=report_progress,
     )
+
+    path_info = get_python_path_info()
+    module_stats = {}  # module -> [num_allocations, total_bytes]
+
     with progress_indicator:
         while True:
             PyErr_CheckSignals()
             ret = reader.nextRecord()
             if ret == RecordResult.RecordResultAllocationRecord:
+                allocation = reader.getLatestAllocation()
                 aggregator.addAllocation(
-                    reader.getLatestAllocation(),
-                    reader.getLatestPythonLocationId(reader.getLatestAllocation()),
+                    allocation,
+                    reader.getLatestPythonLocationId(allocation),
                 )
+
+                if allocation.frame_index != 0:
+                    stack = reader.Py_GetStackFrame(allocation.frame_index)
+                    module_name = get_module_for_stack(stack, path_info)
+                    entry = module_stats.setdefault(module_name, [0, 0])
+                    entry[0] += 1
+                    entry[1] += allocation.size
                 progress_indicator.update(1)
             elif ret == RecordResult.RecordResultMemoryRecord:
                 pass
@@ -1554,6 +1569,22 @@ def compute_statistics(
         for count_and_loc in aggregator.topLocationsByCount(num_largest)
     ]
 
+    top_modules_by_count = sorted(
+        module_stats.items(), key=lambda x: x[1][0], reverse=True
+    )[:num_largest]
+    top_allocations_by_module_by_count = [
+        (name, count, total_bytes)
+        for name, (count, total_bytes) in top_modules_by_count
+    ]
+
+    top_modules_by_size = sorted(
+        module_stats.items(), key=lambda x: x[1][1], reverse=True
+    )[:num_largest]
+    top_allocations_by_module = [
+        (name, count, total_bytes)
+        for name, (count, total_bytes) in top_modules_by_size
+    ]
+
     # And we're done!
     cdef uint64_t peak_memory = aggregator.peakBytesAllocated()
     return Stats(
@@ -1565,6 +1596,8 @@ def compute_statistics(
         allocation_count_by_allocator=allocation_count_by_allocator,
         top_locations_by_size=top_locations_by_size,
         top_locations_by_count=top_locations_by_count,
+        top_allocations_by_module=top_allocations_by_module,
+        top_allocations_by_module_by_count=top_allocations_by_module_by_count,
     )
 
 

src/memray/_stats.py

@@ -1,4 +1,7 @@
 from dataclasses import dataclass
+from dataclasses import field
+from typing import List
+from typing import Tuple
 
 from ._metadata import Metadata
 
@@ -13,3 +16,7 @@ class Stats:
     allocation_count_by_allocator: dict
     top_locations_by_size: list
     top_locations_by_count: list
+    top_allocations_by_module: List[Tuple[str, int, int]] = field(default_factory=list)
+    top_allocations_by_module_by_count: List[Tuple[str, int, int]] = field(
+        default_factory=list
+    )

src/memray/_stats.pyi

@@ -13,3 +13,5 @@ class Stats:
     allocation_count_by_allocator: dict[str, int]
     top_locations_by_size: list[tuple[PythonStackElement, int]]
     top_locations_by_count: list[tuple[PythonStackElement, int]]
+    top_allocations_by_module: list[tuple[str, int, int]]
+    top_allocations_by_module_by_count: list[tuple[str, int, int]]

src/memray/reporters/module_tools.py

@@ -0,0 +1,109 @@
+"""Utilities for extracting module names from file paths."""
+
+import site
+import sys
+import sysconfig
+from pathlib import Path
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from typing_extensions import TypedDict
+
+
+class PathInfo(TypedDict):
+    stdlib: Optional[Path]
+    site_packages: List[Path]
+    sys_path: List[Path]
+
+
+def get_python_path_info() -> PathInfo:
+    """Get information about Python's search paths.
+
+    Returns:
+        dict: Dictionary containing stdlib path, site-packages paths, and sys.path entries.
+    """
+    libdest = sysconfig.get_config_var("LIBDEST")
+    stdlib: Optional[Path] = Path(libdest) if libdest else None
+
+    # Get site-packages directories
+    site_packages: List[Path] = [Path(p) for p in site.getsitepackages()]
+
+    # Get user site-packages
+    user_site = site.getusersitepackages()
+    if Path(user_site).exists():
+        site_packages.append(Path(user_site))
+
+    return {
+        "stdlib": stdlib,
+        "site_packages": site_packages,
+        "sys_path": [Path(p) for p in sys.path if p],
+    }
+
+
+def _is_relative_to(path: Path, base: Path) -> bool:
+    try:
+        path.relative_to(base)
+        return True
+    except ValueError:
+        return False
+
+
+def extract_module_name_and_type(filename: str, path_info: PathInfo) -> Tuple[str, str]:
+    """Extract Python module name and type from file path.
+
+    Returns:
+        tuple: (module_name, module_type) where module_type is one of:
+               'stdlib', 'site-packages', 'project', or 'unknown'
+    """
+    if not filename:
+        return ("unknown", "unknown")
+
+    if filename.startswith("<frozen "):
+        return (filename[len("<frozen ") : -1], "stdlib")
+
+    file_path = Path(filename)
+
+    for site_pkg in path_info["site_packages"]:
+        if _is_relative_to(file_path, site_pkg):
+            return (_path_to_module(file_path.relative_to(site_pkg)), "site-packages")
+
+    if path_info["stdlib"] and _is_relative_to(file_path, path_info["stdlib"]):
+        return (_path_to_module(file_path.relative_to(path_info["stdlib"])), "stdlib")
+
+    for path_entry in path_info["sys_path"]:
+        if _is_relative_to(file_path, path_entry):
+            return (_path_to_module(file_path.relative_to(path_entry)), "project")
+
+    # Fallback: use just the filename, not the full absolute path
+    return (_path_to_module(Path(file_path.name)), "unknown")
+
+
+def _path_to_module(path: Path) -> str:
+    if path.is_absolute():
+        raise ValueError(f"Expected a relative path, got: {path}")
+
+    if path.name == "__init__.py":
+        path = path.parent
+    elif path.suffix == ".py":
+        path = path.with_suffix("")
+
+    return ".".join(path.parts)
+
+
+def get_module_for_stack(
+    stack: Iterable[Tuple[str, str, int]],
+    path_info: PathInfo,
+) -> str:
+    """Find the top-level module of the closest non-stdlib frame in a stack.
+
+    Walks frames from leaf to root, returning the first non-stdlib module's
+    top-level package name. Returns "__main__" if every frame is stdlib
+    or the stack is empty.
+    """
+    for frame in stack:
+        module_name, module_type = extract_module_name_and_type(frame[1], path_info)
+        if module_type != "stdlib":
+            return module_name.split(".")[0]
+    return "__main__"

src/memray/reporters/stats.py

@@ -162,6 +162,25 @@ def _render_to_terminal(self, histogram_params: Dict[str, int]) -> None:
         for location, count in self._get_top_allocations_by_count():
             print(f"\t- {self._format_location(location)} -> {count}")
 
+        for modules, sort_label in [
+            (self._stats.top_allocations_by_module, "by size"),
+            (
+                self._stats.top_allocations_by_module_by_count,
+                "by number of allocations",
+            ),
+        ]:
+            if not modules:
+                continue
+            print()
+            rich.print(
+                f"🥇 [bold]Top {self.num_largest} allocating modules ({sort_label}):[/]"
+            )
+            for module_name, num_allocs, total_bytes in modules:
+                print(
+                    f"\t- {module_name}: {size_fmt(total_bytes)} total"
+                    f" | {num_allocs} allocations"
+                )
+
     def _render_to_json(self, histogram_params: Dict[str, int], out_path: Path) -> None:
         alloc_size_hist = describe_histogram_databins(
             get_histogram_databins(
@@ -189,6 +208,26 @@ def _render_to_json(self, histogram_params: Dict[str, int], out_path: Path) -> N
                 {"location": self._format_location(location), "count": count}
                 for location, count in self._get_top_allocations_by_count()
             ],
+            "top_allocations_by_module": [
+                {
+                    "module": module_name,
+                    "num_allocations": num_allocs,
+                    "total_bytes": total_bytes,
+                }
+                for module_name, num_allocs, total_bytes in (
+                    self._stats.top_allocations_by_module
+                )
+            ],
+            "top_allocations_by_module_by_count": [
+                {
+                    "module": module_name,
+                    "num_allocations": num_allocs,
+                    "total_bytes": total_bytes,
+                }
+                for module_name, num_allocs, total_bytes in (
+                    self._stats.top_allocations_by_module_by_count
+                )
+            ],
             "metadata": metadata,
         }
 

tests/integration/test_tracking.py

@@ -1797,3 +1797,46 @@ def test_temporary_allocations_when_filling_vector_without_preallocating_small_b
         assert record.n_allocations == 1
         assert record.allocator == AllocatorType.MALLOC
         assert record.size == 2 << 10
+
+
+@pytest.mark.parametrize(["allocator_func", "allocator_type"], ALLOCATORS)
+def test_module_stats_only_counts_allocations(allocator_func, allocator_type, tmp_path):
+    # GIVEN
+    output = tmp_path / "test.bin"
+    allocator = MemoryAllocator()
+
+    # WHEN
+    with Tracker(output):
+        res = getattr(allocator, allocator_func)(1234)
+        if res:
+            allocator.free()
+
+    if not res:
+        pytest.skip(f"Allocator {allocator_func} not supported on this platform")
+
+    # THEN
+    stats = compute_statistics(str(output), num_largest=5)
+    total_module_allocs = sum(count for _, count, _ in stats.top_allocations_by_module)
+    assert total_module_allocs <= stats.total_num_allocations
+
+
+def test_module_stats_attributes_to_non_stdlib(tmp_path):
+    # GIVEN
+    output = tmp_path / "test.bin"
+
+    def tracking_function():
+        return [{str(k): k for k in range(50)} for _ in range(1000)]
+
+    # WHEN
+    with Tracker(output, trace_python_allocators=True):
+        tracking_function()
+
+    # THEN
+    stats = compute_statistics(str(output), num_largest=5)
+    module_names = [name for name, _, _ in stats.top_allocations_by_module]
+    assert len(module_names) > 0
+    assert "stdlib" not in module_names
+    for module_name, count, total_bytes in stats.top_allocations_by_module:
+        assert isinstance(module_name, str)
+        assert count > 0
+        assert total_bytes > 0