unified stats function

Author: deadlovelll

Lib/profiling/sampling/binary_reader.py

@@ -103,6 +103,28 @@ def get_stats(self):
         return self._reader.get_stats()
 
 
+def _set_replay_stats(collector, info, sample_interval_usec):
+    """Pass recorded statistics to a collector using the live sampling API."""
+    duration_sec = info.get("duration_sec")
+    sample_rate = info.get("sample_rate")
+    mode = info.get("mode")
+    if (
+        duration_sec is not None
+        and sample_rate is not None
+        and hasattr(collector, "set_stats")
+    ):
+        collector.set_stats(
+            sample_interval_usec,
+            duration_sec,
+            sample_rate,
+            error_rate=info.get("error_rate"),
+            missed_samples=info.get("missed_samples"),
+            mode=mode,
+        )
+    elif hasattr(collector, "set_mode"):
+        collector.set_mode(mode)
+
+
 def convert_binary_to_format(input_file, output_file, output_format,
                              sample_interval_usec=None, progress_callback=None):
     """Convert a binary profiling file to another format.
@@ -137,10 +159,7 @@ def convert_binary_to_format(input_file, output_file, output_format,
 
         # Replay samples through collector
         count = reader.replay_samples(collector, progress_callback)
-        if hasattr(collector, "set_replay_stats"):
-            collector.set_replay_stats(info)
-        if hasattr(collector, "set_mode"):
-            collector.set_mode(info.get("mode"))
+        _set_replay_stats(collector, info, interval)
 
         # Export to target format
         collector.export(output_file)

Lib/profiling/sampling/cli.py

@@ -22,7 +22,7 @@
 from .gecko_collector import GeckoCollector
 from .jsonl_collector import JsonlCollector
 from .binary_collector import BinaryCollector
-from .binary_reader import BinaryReader
+from .binary_reader import BinaryReader, _set_replay_stats
 from .constants import (
     MICROSECONDS_PER_SECOND,
     PROFILING_MODE_ALL,
@@ -793,10 +793,7 @@ def progress_callback(current, total):
             )
 
     count = reader.replay_samples(collector, progress_callback)
-    if hasattr(collector, "set_replay_stats"):
-        collector.set_replay_stats(info)
-    if hasattr(collector, "set_mode"):
-        collector.set_mode(info.get("mode"))
+    _set_replay_stats(collector, info, interval)
     print()
 
     if args.format == "pstats":

Lib/profiling/sampling/heatmap_collector.py

@@ -407,20 +407,6 @@ def set_stats(self, sample_interval_usec, duration_sec, sample_rate, error_rate=
         }
         self.stats.update(kwargs)
 
-    def set_replay_stats(self, info):
-        """Restore measured statistics stored in a binary profile."""
-        duration_sec = info.get("duration_sec")
-        sample_rate = info.get("sample_rate")
-        if duration_sec is None or sample_rate is None:
-            return
-        self.set_stats(
-            self.sample_interval_usec,
-            duration_sec,
-            sample_rate,
-            error_rate=info.get("error_rate"),
-            missed_samples=info.get("missed_samples"),
-        )
-
     def process_frames(self, frames, thread_id, weight=1):
         """Process stack frames and count samples per line.
 

Lib/profiling/sampling/stack_collector.py

@@ -152,20 +152,6 @@ def set_stats(self, sample_interval_usec, duration_sec, sample_rate,
             "mode": mode
         }
 
-    def set_replay_stats(self, info):
-        """Restore measured statistics stored in a binary profile."""
-        duration_sec = info.get("duration_sec")
-        sample_rate = info.get("sample_rate")
-        if duration_sec is None or sample_rate is None:
-            return
-        self.set_stats(
-            self.sample_interval_usec,
-            duration_sec,
-            sample_rate,
-            error_rate=info.get("error_rate"),
-            missed_samples=info.get("missed_samples"),
-            mode=self.stats.get("mode"),
-        )
     def set_mode(self, mode):
         self.stats["mode"] = mode
 

Lib/test/test_profiling/test_sampling_profiler/test_collectors.py

@@ -13,6 +13,7 @@
 try:
     import _remote_debugging  # noqa: F401
     from profiling.sampling import gecko_collector
+    from profiling.sampling.binary_reader import _set_replay_stats
     from profiling.sampling.pstats_collector import PstatsCollector
     from profiling.sampling.stack_collector import (
         CollapsedStackCollector,
@@ -1218,12 +1219,12 @@ def test_flamegraph_collector_stats_accumulation(self):
     def test_flamegraph_collector_restores_replay_stats(self):
         """Replay restores measured values from binary metadata."""
         collector = FlamegraphCollector(1000)
-        collector.set_replay_stats({
+        _set_replay_stats(collector, {
             "duration_sec": 1.25,
             "sample_rate": 4.0,
             "error_rate": 2.5,
             "missed_samples": 1.5,
-        })
+        }, 1000)
 
         self.assertEqual(collector.stats["duration_sec"], 1.25)
         self.assertEqual(collector.stats["sample_rate"], 4.0)
@@ -1233,10 +1234,12 @@ def test_flamegraph_collector_restores_replay_stats(self):
     def test_flamegraph_collector_leaves_legacy_replay_stats_unavailable(self):
         collector = FlamegraphCollector(1000)
         original_stats = collector.stats.copy()
-        collector.set_replay_stats({
+        original_stats["mode"] = PROFILING_MODE_CPU
+        _set_replay_stats(collector, {
             "duration_sec": None,
             "sample_rate": None,
-        })
+            "mode": PROFILING_MODE_CPU,
+        }, 1000)
         self.assertEqual(collector.stats, original_stats)
 
     def test_flamegraph_collector_per_thread_stats(self):