stats: add output_file option to output the stats to a file

Author: AlexTereshenkov

src/python/pants/goal/stats_aggregator.py

@@ -7,6 +7,7 @@
 import logging
 from collections import Counter
 from dataclasses import dataclass
+from typing import Optional
 
 from pants.engine.internals.scheduler import Workunit
 from pants.engine.rules import collect_rules, rule
@@ -17,7 +18,7 @@
     WorkunitsCallbackFactoryRequest,
 )
 from pants.engine.unions import UnionRule
-from pants.option.option_types import BoolOption
+from pants.option.option_types import BoolOption, StrOption
 from pants.option.subsystem import Subsystem
 from pants.util.collections import deep_getsizeof
 from pants.util.strutil import softwrap
@@ -55,13 +56,30 @@ class StatsAggregatorSubsystem(Subsystem):
         ),
         advanced=True,
     )
+    output_file = StrOption(
+        default=None,
+        metavar="<path>",
+        help="Output the stats to this file. If unspecified, outputs to stdout.",
+    )
+
+
+def log_or_write_to_file(output_file: Optional[str], text: str, logger: logging.Logger) -> None:
+    """Send text to the stdout or write to the output file."""
+    if output_file:
+        with open(output_file, "w") as fh:
+            fh.write(text)
+    else:
+        logger.info(text)
 
 
 class StatsAggregatorCallback(WorkunitsCallback):
-    def __init__(self, *, log: bool, memory: bool, has_histogram_module: bool) -> None:
+    def __init__(
+        self, *, log: bool, memory: bool, output_file: Optional[str], has_histogram_module: bool
+    ) -> None:
         super().__init__()
         self.log = log
         self.memory = memory
+        self.output_file = output_file
         self.has_histogram_module = has_histogram_module
 
     @property
@@ -80,6 +98,8 @@ def __call__(
         if not finished:
             return
 
+        output_contents = ""
+
         if self.log:
             # Capture global counters.
             counters = Counter(context.get_metrics())
@@ -93,7 +113,7 @@ def __call__(
             counter_lines = "\n".join(
                 f"  {name}: {count}" for name, count in sorted(counters.items())
             )
-            logger.info(f"Counters:\n{counter_lines}")
+            output_contents += f"Counters:\n{counter_lines}"
 
         if self.memory:
             ids: set[int] = set()
@@ -115,18 +135,23 @@ def __call__(
             memory_lines = "\n".join(
                 f"  {size}\t\t{count}\t\t{name}" for size, count, name in sorted(entries)
             )
-            logger.info(f"Memory summary (total size in bytes, count, name):\n{memory_lines}")
+            output_contents += (
+                f"\nMemory summary (total size in bytes, count, name):\n{memory_lines}"
+            )
 
         if not (self.log and self.has_histogram_module):
+            log_or_write_to_file(self.output_file, output_contents, logger)
             return
+
         from hdrh.histogram import HdrHistogram  # pants: no-infer-dep
 
         histograms = context.get_observation_histograms()["histograms"]
         if not histograms:
-            logger.info("No observation histogram were recorded.")
+            output_contents += "\nNo observation histogram were recorded."
+            log_or_write_to_file(self.output_file, output_contents, logger)
             return
 
-        logger.info("Observation histogram summaries:")
+        output_contents += "\nObservation histogram summaries:"
         for name, encoded_histogram in histograms.items():
             # Note: The Python library for HDR Histogram will only decode compressed histograms
             # that are further encoded with base64. See
@@ -138,8 +163,8 @@ def __call__(
                     [25, 50, 75, 90, 95, 99]
                 ).items()
             )
-            logger.info(
-                f"Summary of `{name}` observation histogram:\n"
+            output_contents += (
+                f"\nSummary of `{name}` observation histogram:\n"
                 f"  min: {histogram.get_min_value()}\n"
                 f"  max: {histogram.get_max_value()}\n"
                 f"  mean: {histogram.get_mean_value():.3f}\n"
@@ -148,6 +173,7 @@ def __call__(
                 f"  sum: {int(histogram.get_mean_value() * histogram.total_count)}\n"
                 f"{percentile_to_vals}"
             )
+        log_or_write_to_file(self.output_file, output_contents, logger)
 
 
 @dataclass(frozen=True)
@@ -178,6 +204,7 @@ def construct_callback(
             StatsAggregatorCallback(
                 log=subsystem.log,
                 memory=subsystem.memory_summary,
+                output_file=subsystem.output_file,
                 has_histogram_module=has_histogram_module,
             )
             if subsystem.log or subsystem.memory_summary

src/python/pants/goal/stats_aggregator_integration_test.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import re
+from pathlib import Path
 
 from pants.testutil.pants_integration_test import run_pants, setup_tmpdir
 
@@ -46,3 +47,18 @@ def test_warn_if_no_histograms() -> None:
     assert "Counters:" in result.stderr
     assert "Please run with `--plugins=hdrhistogram`" in result.stderr
     assert "Observation histogram summaries:" not in result.stderr
+
+
+def test_writing_to_output_file() -> None:
+    with setup_tmpdir({"src/py/app.py": "print(0)\n", "src/py/BUILD": "python_sources()"}):
+        argv = [
+            "--backend-packages=['pants.backend.python']",
+            "--stats-log",
+            "--stats-memory-summary",
+            "--stats-output-file=stats.txt",
+            "roots",
+        ]
+        run_pants(argv).assert_success()
+        output_file_contents = Path("stats.txt").read_text()
+        for item in ("Counters:", "Memory summary"):
+            assert item in output_file_contents