Add stats --json output

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
 -----------
 
@@ -34,6 +36,18 @@ previously generated using :doc:`the run subcommand <run>`.
 
 The output will be printed directly to the standard output of the terminal.
 
+JSON Output
+-----------
+
+If you supply the ``--json`` flag, the ``stats`` subcommand will write its
+output to a JSON file, rather than to the terminal. Like other commands that
+output to files, the default output file name is based on the name of your
+capture file, but it can be overridden with the ``-o`` / ``--output`` option.
+By default Memray will refuse to overwrite an existing file, but you can force
+it to by supplying the ``-f`` / ``--force`` option.
+
+Note that new fields may be added to the JSON output over time, though we'll
+try to avoid removing existing fields.
 
 CLI Reference
 -------------

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

@@ -1,6 +1,7 @@
 import argparse
 import os
 from pathlib import Path
+from typing import Optional
 
 from memray._errors import MemrayCommandError
 from memray._memray import compute_statistics
@@ -33,6 +34,26 @@ 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,
+        )
+        parser.add_argument(
+            "-o",
+            "--output",
+            help="Output file name for JSON output",
+            default=None,
+        )
+        parser.add_argument(
+            "-f",
+            "--force",
+            help="If the JSON output file already exists, overwrite it",
+            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():
@@ -49,5 +70,24 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
                 exit_code=1,
             )
 
+        json_output_file: Optional[Path] = None
+        if args.json:
+            if args.output:
+                json_output_file = Path(args.output)
+            else:
+                filename = str(result_path.name) + ".json"
+                if filename.startswith("memray-"):
+                    filename = filename[len("memray-") :]
+                filename = "memray-stats-" + filename
+                json_output_file = result_path.with_name(filename)
+
+            if not args.force and json_output_file.exists():
+                raise MemrayCommandError(
+                    f"File already exists, will not overwrite: {json_output_file}",
+                    exit_code=1,
+                )
+
         reporter = StatsReporter(stats, args.num_largest)
-        reporter.render()
+        reporter.render(json_output_file=json_output_file)
+        if json_output_file is not None:
+            print(f"Wrote {json_output_file}")

src/memray/reporters/stats.py

@@ -1,15 +1,23 @@
+import datetime
+import json
 import math
 from collections import Counter
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any
 from typing import Dict
 from typing import Iterator
 from typing import List
+from typing import Optional
 from typing import Tuple
 
 import rich
 
 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:
@@ -30,6 +38,19 @@ def get_histogram_databins(data: Dict[int, int], bins: int) -> List[Tuple[int, i
     return [(steps[b], dist[b]) for b in range(bins)]
 
 
+def describe_histogram_databins(
+    databins: List[Tuple[int, int]]
+) -> List[Dict[str, int]]:
+    ret: List[Dict[str, int]] = []
+    start = 0
+    for i, (end, count) in enumerate(databins):
+        # The max size for the last bucket is inclusive, not exclusive
+        adjustment = 1 if i != len(databins) - 1 else 0
+        ret.append(dict(min_bytes=start, max_bytes=end - adjustment, count=count))
+        start = end
+    return ret
+
+
 def draw_histogram(
     data: Dict[int, int], bins: int, *, hist_scale_factor: int = 25
 ) -> str:
@@ -89,7 +110,17 @@ 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, json_output_file: Optional[Path] = None) -> None:
+        histogram_params = dict(
+            num_bins=10,
+            histogram_scale_factor=25,
+        )
+        if json_output_file:
+            self._render_to_json(histogram_params, json_output_file)
+        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 +129,68 @@ 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 = describe_histogram_databins(
+            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[str, Any] = 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 +199,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)