stats: Support --force and --output for JSON mode

Signed-off-by: Matt Wozniski <mwozniski@bloomberg.net>

Author: godlygeek

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
@@ -39,6 +40,19 @@ def valid_positive_int(value: str) -> int:
             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)
@@ -56,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(to_json=args.json, result_path=result_path)
+        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,14 +1,14 @@
 import datetime
 import json
 import math
-import re
 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
@@ -110,17 +110,13 @@ 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, to_json: bool = False, result_path: Path = Path(".")) -> None:
+    def render(self, json_output_file: Optional[Path] = None) -> 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)
+        if json_output_file:
+            self._render_to_json(histogram_params, json_output_file)
         else:
             self._render_to_terminal(histogram_params)
 

tests/integration/test_main.py

@@ -1,4 +1,5 @@
 import contextlib
+import json
 import os
 import platform
 import pty
@@ -863,6 +864,138 @@ def test_report_generated(self, tmp_path, simple_test_file):
         # THEN
         assert "VALLOC" in output
 
+    def test_json_generated(self, tmp_path, simple_test_file):
+        # GIVEN
+        results_file, _ = generate_sample_results(tmp_path, simple_test_file)
+        json_file = tmp_path / "memray-stats-result.bin.json"
+
+        # WHEN
+        subprocess.check_output(
+            [
+                sys.executable,
+                "-m",
+                "memray",
+                "stats",
+                "--json",
+                str(results_file),
+            ],
+            cwd=str(tmp_path),
+            text=True,
+        )
+
+        # THEN
+        assert json_file.exists()
+        assert isinstance(json.loads(json_file.read_text()), dict)
+
+    def test_json_generated_to_pretty_file_name(self, tmp_path, simple_test_file):
+        # GIVEN
+        orig_results_file, _ = generate_sample_results(tmp_path, simple_test_file)
+        results_file = orig_results_file.with_name("memray-foobar.bin")
+        orig_results_file.rename(results_file)
+        json_file = tmp_path / "memray-stats-foobar.bin.json"
+
+        # WHEN
+        subprocess.check_output(
+            [
+                sys.executable,
+                "-m",
+                "memray",
+                "stats",
+                "--json",
+                str(results_file),
+            ],
+            cwd=str(tmp_path),
+            text=True,
+        )
+
+        # THEN
+        assert json_file.exists()
+        assert isinstance(json.loads(json_file.read_text()), dict)
+
+    def test_json_generated_to_known_file(self, tmp_path, simple_test_file):
+        # GIVEN
+        results_file, _ = generate_sample_results(tmp_path, simple_test_file)
+        json_file = tmp_path / "output.json"
+
+        # WHEN
+        subprocess.check_output(
+            [
+                sys.executable,
+                "-m",
+                "memray",
+                "stats",
+                "--json",
+                "-o",
+                str(json_file),
+                str(results_file),
+            ],
+            cwd=str(tmp_path),
+            text=True,
+        )
+
+        # THEN
+        assert json_file.exists()
+        assert isinstance(json.loads(json_file.read_text()), dict)
+
+    def test_json_generated_to_existing_known_file(self, tmp_path, simple_test_file):
+        # GIVEN
+        results_file, _ = generate_sample_results(tmp_path, simple_test_file)
+        json_file = tmp_path / "output.json"
+        json_file.write_text("oops")
+
+        # WHEN
+        try:
+            exc = None
+            subprocess.check_output(
+                [
+                    sys.executable,
+                    "-m",
+                    "memray",
+                    "stats",
+                    "--json",
+                    "-o",
+                    str(json_file),
+                    str(results_file),
+                ],
+                cwd=str(tmp_path),
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+        except subprocess.CalledProcessError as e:
+            exc = e
+
+        # THEN
+        assert exc is not None
+        assert "File already exists, will not overwrite" in exc.stderr
+
+    def test_json_overwrites_existing_known_file(self, tmp_path, simple_test_file):
+        # GIVEN
+        results_file, _ = generate_sample_results(tmp_path, simple_test_file)
+        json_file = tmp_path / "output.json"
+        json_file.write_text("oops")
+
+        # WHEN
+        subprocess.check_output(
+            [
+                sys.executable,
+                "-m",
+                "memray",
+                "stats",
+                "--json",
+                "--force",
+                "--output",
+                str(json_file),
+                str(results_file),
+            ],
+            cwd=str(tmp_path),
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+
+        # THEN
+        assert json_file.exists()
+        assert isinstance(json.loads(json_file.read_text()), dict)
+
     def test_report_detects_corrupt_input(self, tmp_path):
         # GIVEN
         bad_file = Path(tmp_path) / "badfile.bin"

tests/unit/test_stats_reporter.py

@@ -1,3 +1,4 @@
+import json
 from collections import Counter
 from datetime import datetime
 from typing import List
@@ -391,9 +392,9 @@ def test_stats_output(fake_stats):
 
 
 def test_stats_output_json(fake_stats, tmp_path):
+    output_file = tmp_path / "json.out"
     reporter = StatsReporter(fake_stats, 5)
-    with patch("json.dump") as json_dump:
-        reporter.render(to_json=True, result_path=tmp_path)
+    reporter.render(json_output_file=output_file)
     expected = {
         "total_num_allocations": 20,
         "total_bytes_allocated": 3341500,
@@ -437,5 +438,5 @@ def test_stats_output_json(fake_stats, tmp_path):
             "has_native_traces": False,
         },
     }
-    actual = json_dump.call_args[0][0]
+    actual = json.loads(output_file.read_text())
     assert expected == actual