Allow memray stats to output a JSON report

Signed-off-by: Hao Sun <helioeccentric@gmail.com>

Author: oddsun

docs/stats.rst

@@ -20,6 +20,8 @@ The output includes the following:
 
 * Stack trace and **count** of the top 'n' largest allocating locations by number of allocations (*default: 5*, configurable with the ``-n`` command line param)
 
+* (for JSON output only) Metadata about the tracked process
+
 Basic Usage
 -----------
 

news/377.feature.rst

@@ -0,0 +1 @@
+Allow ``memray stats`` to output a JSON report via ``--json`` flag.

src/memray/_stats.pyi

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
 
-from ._memray import AllocatorType
 from ._memray import PythonStackElement
 from ._metadata import Metadata
 
@@ -11,6 +10,6 @@ class Stats:
     total_memory_allocated: int
     peak_memory_allocated: int
     allocation_count_by_size: dict[int, int]
-    allocation_count_by_allocator: dict[AllocatorType, int]
+    allocation_count_by_allocator: dict[str, int]
     top_locations_by_size: list[tuple[PythonStackElement, int]]
     top_locations_by_count: list[tuple[PythonStackElement, int]]

src/memray/commands/stats.py

@@ -33,6 +33,13 @@ def valid_positive_int(value: str) -> int:
             default=5,
         )
 
+        parser.add_argument(
+            "--json",
+            help="Exports stats to a JSON file",
+            action="store_true",
+            default=False,
+        )
+
     def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
         result_path = Path(args.results)
         if not result_path.exists() or not result_path.is_file():
@@ -50,4 +57,4 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
             )
 
         reporter = StatsReporter(stats, args.num_largest)
-        reporter.render()
+        reporter.render(to_json=args.json, result_path=result_path)

src/memray/reporters/stats.py

@@ -1,5 +1,10 @@
+import datetime
+import json
 import math
+import re
 from collections import Counter
+from dataclasses import asdict
+from pathlib import Path
 from typing import Dict
 from typing import Iterator
 from typing import List
@@ -10,6 +15,8 @@
 from memray._memray import size_fmt
 from memray._stats import Stats
 
+PythonStackElement = Tuple[str, str, int]
+
 
 def get_histogram_databins(data: Dict[int, int], bins: int) -> List[Tuple[int, int]]:
     if bins <= 0:
@@ -89,7 +96,21 @@ def __init__(self, stats: Stats, num_largest: int):
             raise ValueError(f"Invalid input num_largest={num_largest}, should be >=1")
         self.num_largest = num_largest
 
-    def render(self) -> None:
+    def render(self, to_json: bool = False, result_path: Path = Path(".")) -> None:
+        histogram_params = dict(
+            num_bins=10,
+            histogram_scale_factor=25,
+        )
+        if to_json:
+            # appends suffix (works better with --follow-fork)
+            out_name = result_path.with_suffix(result_path.suffix + ".json").name
+            out_name = re.sub("^memray-", "", out_name)
+            out_path = result_path.parent / f"memray-stats-{out_name}"
+            self._render_to_json(histogram_params, out_path)
+        else:
+            self._render_to_terminal(histogram_params)
+
+    def _render_to_terminal(self, histogram_params: Dict[str, int]) -> None:
         rich.print("📏 [bold]Total allocations:[/]")
         print(f"\t{self._stats.total_num_allocations}")
 
@@ -98,35 +119,66 @@ def render(self) -> None:
         print(f"\t{size_fmt(self._stats.total_memory_allocated)}")
 
         print()
-        num_bins = 10
-        histogram_scale_factor = 25
+
         rich.print("📊 [bold]Histogram of allocation size:[/]")
         histogram = draw_histogram(
             self._stats.allocation_count_by_size,
-            num_bins,
-            hist_scale_factor=histogram_scale_factor,
+            histogram_params["num_bins"],
+            hist_scale_factor=histogram_params["histogram_scale_factor"],
         )
         print(f"\t{histogram}")
 
         print()
         rich.print("📂 [bold]Allocator type distribution:[/]")
-        for entry in self._get_allocator_type_distribution():
-            print(f"\t {entry}")
+        for allocator_name, count in self._get_allocator_type_distribution():
+            print(f"\t {allocator_name}: {count}")
 
         print()
         rich.print(
             f"🥇 [bold]Top {self.num_largest} largest allocating locations (by size):[/]"
         )
-        for entry in self._get_top_allocations_by_size():
-            print(f"\t- {entry}")
+        for location, size in self._get_top_allocations_by_size():
+            print(f"\t- {self._format_location(location)} -> {size_fmt(size)}")
 
         print()
         rich.print(
             f"🥇 [bold]Top {self.num_largest} largest allocating "
             "locations (by number of allocations):[/]"
         )
-        for entry in self._get_top_allocations_by_count():
-            print(f"\t- {entry}")
+        for location, count in self._get_top_allocations_by_count():
+            print(f"\t- {self._format_location(location)} -> {count}")
+
+    def _render_to_json(self, histogram_params: Dict[str, int], out_path: Path) -> None:
+        alloc_size_hist = get_histogram_databins(
+            self._stats.allocation_count_by_size, bins=histogram_params["num_bins"]
+        )
+
+        metadata = asdict(self._stats.metadata)
+        for name, val in metadata.items():
+            if isinstance(val, datetime.datetime):
+                metadata[name] = str(val)
+
+        data = dict(
+            total_num_allocations=self._stats.total_num_allocations,
+            total_bytes_allocated=self._stats.total_memory_allocated,
+            allocation_size_histogram=alloc_size_hist,
+            allocator_type_distribution=[
+                (allocation_type, count)
+                for allocation_type, count in self._get_allocator_type_distribution()
+            ],
+            top_allocations_by_size=[
+                {"location": self._format_location(location), "size": size}
+                for location, size in self._get_top_allocations_by_size()
+            ],
+            top_allocations_by_count=[
+                {"location": self._format_location(location), "count": count}
+                for location, count in self._get_top_allocations_by_count()
+            ],
+            metadata=metadata,
+        )
+
+        with open(out_path, "w") as f:
+            json.dump(data, f, indent=2)
 
     @staticmethod
     def _format_location(loc: Tuple[str, str, int]) -> str:
@@ -135,18 +187,18 @@ def _format_location(loc: Tuple[str, str, int]) -> str:
             return "<stack trace unavailable>"
         return f"{function}:{file}:{line}"
 
-    def _get_top_allocations_by_size(self) -> Iterator[str]:
+    def _get_top_allocations_by_size(self) -> Iterator[Tuple[PythonStackElement, int]]:
         for location, size in self._stats.top_locations_by_size:
-            yield f"{self._format_location(location)} -> {size_fmt(size)}"
+            yield (location, size)
 
-    def _get_top_allocations_by_count(self) -> Iterator[str]:
+    def _get_top_allocations_by_count(self) -> Iterator[Tuple[PythonStackElement, int]]:
         for location, count in self._stats.top_locations_by_count:
-            yield f"{self._format_location(location)} -> {count}"
+            yield (location, count)
 
-    def _get_allocator_type_distribution(self) -> Iterator[str]:
+    def _get_allocator_type_distribution(self) -> Iterator[Tuple[str, int]]:
         for allocator_name, count in sorted(
             self._stats.allocation_count_by_allocator.items(),
             key=lambda item: item[1],
             reverse=True,
         ):
-            yield f"{allocator_name}: {count}"
+            yield (allocator_name, count)

tests/unit/test_stats_reporter.py

@@ -1,11 +1,16 @@
 from collections import Counter
+from datetime import datetime
 from typing import List
 from typing import Optional
 from typing import Tuple
+from unittest.mock import patch
 
 import pytest
 
 from memray import AllocatorType as AT
+from memray._metadata import Metadata
+from memray._stats import Stats
+from memray.reporters.stats import StatsReporter
 from memray.reporters.stats import draw_histogram
 from memray.reporters.stats import get_histogram_databins
 from tests.utils import MockAllocationRecord
@@ -62,6 +67,62 @@ def _generate_mock_allocations(
     return snapshot
 
 
+# data generator for tests
+@pytest.fixture(scope="module")
+def fake_stats():
+    mem_allocation_list = [
+        2500,
+        11000,
+        11000,
+        12000,
+        60000,
+        65000,
+        120000,
+        125000,
+        125000,
+        160000,
+        170000,
+        180000,
+        800000,
+        1500000,
+    ]
+
+    s = Stats(
+        metadata=Metadata(
+            start_time=datetime(2023, 1, 1, 1),
+            end_time=datetime(2023, 1, 1, 2),
+            total_allocations=sum(mem_allocation_list),
+            total_frames=10,
+            peak_memory=max(mem_allocation_list),
+            command_line="fake stats",
+            pid=123456,
+            python_allocator="pymalloc",
+            has_native_traces=False,
+        ),
+        total_num_allocations=20,
+        total_memory_allocated=sum(mem_allocation_list),
+        peak_memory_allocated=max(mem_allocation_list),
+        allocation_count_by_size=Counter(mem_allocation_list),
+        allocation_count_by_allocator={
+            AT.MALLOC.name: 1013,
+            AT.REALLOC.name: 797,
+            AT.CALLOC.name: 152,
+            AT.MMAP.name: 4,
+        },
+        top_locations_by_count=[
+            (("fake_func", "fake.py", 5), 20),
+            (("fake_func2", "fake.py", 10), 50),
+            (("__main__", "fake.py", 15), 1),
+        ],
+        top_locations_by_size=[
+            (("fake_func", "fake.py", 5), 5 * 2**20),
+            (("fake_func2", "fake.py", 10), 3 * 2**10),
+            (("__main__", "fake.py", 15), 4),
+        ],
+    )
+    return s
+
+
 # tests begin here
 def test_get_histogram_databins():
     # GIVEN
@@ -279,3 +340,102 @@ def test_draw_histogram_invalid_input():
     # test#3 - Invalid hist_scale_factor value
     with pytest.raises(ValueError):
         _ = draw_histogram([100, 200, 300], bins=5, hist_scale_factor=0)
+
+
+def test_stats_output(fake_stats):
+    reporter = StatsReporter(fake_stats, 5)
+    with patch("builtins.print") as mocked_print:
+        with patch("rich.print", print):
+            reporter.render()
+    expected = (
+        "📏 [bold]Total allocations:[/]\n"
+        "\t20\n"
+        "\n"
+        "📦 [bold]Total memory allocated:[/]\n"
+        "\t3.187MB\n"
+        "\n"
+        "📊 [bold]Histogram of allocation size:[/]\n"
+        "\tmin: 2.441KB\n"
+        "\t----------------------------------------\n"
+        "\t< 4.628KB  : 1 ▇▇▇▇▇\n"
+        "\t< 8.775KB  : 0 \n"
+        "\t< 16.637KB : 3 ▇▇▇▇▇▇▇▇▇▇▇▇▇\n"
+        "\t< 31.542KB : 0 \n"
+        "\t< 59.802KB : 1 ▇▇▇▇▇\n"
+        "\t< 113.378KB: 1 ▇▇▇▇▇\n"
+        "\t< 214.954KB: 6 ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇\n"
+        "\t< 407.531KB: 0 \n"
+        "\t< 772.638KB: 0 \n"
+        "\t<=1.431MB  : 2 ▇▇▇▇▇▇▇▇▇\n"
+        "\t----------------------------------------\n"
+        "\tmax: 1.431MB\n"
+        "\n"
+        "📂 [bold]Allocator type distribution:[/]\n"
+        "\t MALLOC: 1013\n"
+        "\t REALLOC: 797\n"
+        "\t CALLOC: 152\n"
+        "\t MMAP: 4\n"
+        "\n"
+        "🥇 [bold]Top 5 largest allocating locations (by size):[/]\n"
+        "\t- fake_func:fake.py:5 -> 5.000MB\n"
+        "\t- fake_func2:fake.py:10 -> 3.000KB\n"
+        "\t- __main__:fake.py:15 -> 4.000B\n"
+        "\n"
+        "🥇 [bold]Top 5 largest allocating locations (by number of allocations):[/]\n"
+        "\t- fake_func:fake.py:5 -> 20\n"
+        "\t- fake_func2:fake.py:10 -> 50\n"
+        "\t- __main__:fake.py:15 -> 1"
+    )
+    printed = "\n".join(" ".join(x[0]) for x in mocked_print.call_args_list)
+    assert expected == printed
+
+
+def test_stats_output_json(fake_stats, tmp_path):
+    reporter = StatsReporter(fake_stats, 5)
+    with patch("json.dump") as json_dump:
+        reporter.render(to_json=True, result_path=tmp_path)
+    expected = {
+        "total_num_allocations": 20,
+        "total_bytes_allocated": 3341500,
+        "allocation_size_histogram": [
+            (4739, 1),
+            (8986, 0),
+            (17036, 3),
+            (32299, 0),
+            (61237, 1),
+            (116099, 1),
+            (220113, 6),
+            (417312, 0),
+            (791181, 0),
+            (1500000, 2),
+        ],
+        "allocator_type_distribution": [
+            ("MALLOC", 1013),
+            ("REALLOC", 797),
+            ("CALLOC", 152),
+            ("MMAP", 4),
+        ],
+        "top_allocations_by_size": [
+            {"location": "fake_func:fake.py:5", "size": 5242880},
+            {"location": "fake_func2:fake.py:10", "size": 3072},
+            {"location": "__main__:fake.py:15", "size": 4},
+        ],
+        "top_allocations_by_count": [
+            {"location": "fake_func:fake.py:5", "count": 20},
+            {"location": "fake_func2:fake.py:10", "count": 50},
+            {"location": "__main__:fake.py:15", "count": 1},
+        ],
+        "metadata": {
+            "start_time": "2023-01-01 01:00:00",
+            "end_time": "2023-01-01 02:00:00",
+            "total_allocations": 3341500,
+            "total_frames": 10,
+            "peak_memory": 1500000,
+            "command_line": "fake stats",
+            "pid": 123456,
+            "python_allocator": "pymalloc",
+            "has_native_traces": False,
+        },
+    }
+    actual = json_dump.call_args[0][0]
+    assert expected == actual