Address feedback and ensuring the TUI works at the end

Author: pablogsal

Lib/profiling/sampling/live_collector/collector.py

@@ -159,6 +159,8 @@ def __init__(
         self.filter_input_mode = False  # Currently entering filter text
         self.filter_input_buffer = ""  # Buffer for filter input
         self.finished = False  # Program has finished, showing final state
+        self.finish_timestamp = None  # When profiling finished (for time freezing)
+        self.finish_wall_time = None  # Wall clock time when profiling finished
 
         # Thread tracking state
         self.thread_ids = []  # List of thread IDs seen
@@ -183,6 +185,20 @@ def __init__(
         # Trend tracking (initialized after colors are set up)
         self._trend_tracker = None
 
+    @property
+    def elapsed_time(self):
+        """Get the elapsed time, frozen when finished."""
+        if self.finished and self.finish_timestamp is not None:
+            return self.finish_timestamp - self.start_time
+        return time.perf_counter() - self.start_time if self.start_time else 0
+
+    @property
+    def current_time_display(self):
+        """Get the current time for display, frozen when finished."""
+        if self.finished and self.finish_wall_time is not None:
+            return time.strftime("%H:%M:%S", time.localtime(self.finish_wall_time))
+        return time.strftime("%H:%M:%S")
+
     def _get_or_create_thread_data(self, thread_id):
         """Get or create ThreadData for a thread ID."""
         if thread_id not in self.per_thread_data:
@@ -384,9 +400,7 @@ def collect(self, stack_frames):
 
     def _prepare_display_data(self, height):
         """Prepare data for display rendering."""
-        elapsed = (
-            time.perf_counter() - self.start_time if self.start_time else 0
-        )
+        elapsed = self.elapsed_time
         stats_list = self._build_stats_list()
 
         # Calculate available space for stats
@@ -707,16 +721,28 @@ def reset_stats(self):
         # Clear trend tracking
         if self._trend_tracker is not None:
             self._trend_tracker.clear()
+        # Reset finished state and finish timestamp
+        self.finished = False
+        self.finish_timestamp = None
+        self.finish_wall_time = None
         self.start_time = time.perf_counter()
         self._last_display_update = self.start_time
 
     def mark_finished(self):
         """Mark the profiling session as finished."""
         self.finished = True
+        # Capture the finish timestamp to freeze all timing displays
+        self.finish_timestamp = time.perf_counter()
+        self.finish_wall_time = time.time()  # Wall clock time for display
         # Force a final display update to show the finished message
         if self.display is not None:
             self._update_display()
 
+    def _handle_finished_input_update(self, had_input):
+        """Update display after input when program is finished."""
+        if self.finished and had_input and self.display is not None:
+            self._update_display()
+
     def _show_terminal_too_small(self, height, width):
         """Display a message when terminal is too small."""
         A_BOLD = self.display.get_attr("A_BOLD")
@@ -809,12 +835,7 @@ def _handle_input(self):
         self.display.set_nodelay(True)
         ch = self.display.get_input()
 
-        # If showing help, any key closes it
-        if self.show_help and ch != -1:
-            self.show_help = False
-            return
-
-        # Handle filter input mode
+        # Handle filter input mode FIRST - takes precedence over all commands
         if self.filter_input_mode:
             if ch == 27:  # ESC key
                 self.filter_input_mode = False
@@ -832,14 +853,19 @@ def _handle_input(self):
                     self.filter_input_buffer = self.filter_input_buffer[:-1]
             elif ch >= 32 and ch < 127:  # Printable characters
                 self.filter_input_buffer += chr(ch)
-            return
 
-        # If finished, only allow 'q' to quit
-        if self.finished:
-            if ch == ord("q") or ch == ord("Q"):
-                self.running = False
+            # Update display if input was processed while finished
+            self._handle_finished_input_update(ch != -1)
             return
 
+        # Handle help toggle keys
+        if ch == ord("h") or ch == ord("H") or ch == ord("?"):
+            self.show_help = not self.show_help
+
+        # If showing help, any other key closes it
+        elif self.show_help and ch != -1:
+            self.show_help = False
+
         # Handle regular commands
         if ch == ord("q") or ch == ord("Q"):
             self.running = False
@@ -850,14 +876,13 @@ def _handle_input(self):
         elif ch == ord("S"):
             self._cycle_sort(reverse=True)
 
-        elif ch == ord("h") or ch == ord("H") or ch == ord("?"):
-            self.show_help = not self.show_help
-
         elif ch == ord("p") or ch == ord("P"):
             self.paused = not self.paused
 
         elif ch == ord("r") or ch == ord("R"):
-            self.reset_stats()
+            # Don't allow reset when profiling is finished
+            if not self.finished:
+                self.reset_stats()
 
         elif ch == ord("+") or ch == ord("="):
             # Decrease update interval (faster refresh)
@@ -915,6 +940,9 @@ def _handle_input(self):
                         self.current_thread_index + 1
                     ) % len(self.thread_ids)
 
+        # Update display if input was processed while finished
+        self._handle_finished_input_update(ch != -1)
+
     def init_curses(self, stdscr):
         """Initialize curses display and suppress stdout/stderr."""
         self.stdscr = stdscr

Lib/profiling/sampling/live_collector/widgets.py

@@ -175,7 +175,7 @@ def draw_header_info(self, line, width, elapsed):
         self.add_str(line, 0, title, A_BOLD | self.colors["cyan"])
         line += 1
 
-        current_time = time.strftime("%H:%M:%S")
+        current_time = self.collector.current_time_display
         uptime = self.format_uptime(elapsed)
 
         # Calculate display refresh rate

Lib/test/test_profiling/test_sampling_profiler/test_live_collector_interaction.py

@@ -318,27 +318,30 @@ def test_finished_state_displays_banner(self):
         self.assertTrue(self.display.contains_text("PROFILING COMPLETE"))
         self.assertTrue(self.display.contains_text("Press 'q' to Quit"))
 
-    def test_finished_state_ignores_most_input(self):
-        """Test that finished state only responds to 'q' key."""
+    def test_finished_state_allows_ui_controls(self):
+        """Test that finished state allows UI controls but prioritizes quit."""
         self.collector.finished = True
         self.collector.running = True
 
-        # Try pressing 's' (sort) - should be ignored
+        # Try pressing 's' (sort) - should work and trigger display update
+        original_sort = self.collector.sort_by
         self.display.simulate_input(ord("s"))
         self.collector._handle_input()
         self.assertTrue(self.collector.running)  # Still running
+        self.assertNotEqual(self.collector.sort_by, original_sort)  # Sort changed
 
-        # Try pressing 'p' (pause) - should be ignored
+        # Try pressing 'p' (pause) - should work
         self.display.simulate_input(ord("p"))
         self.collector._handle_input()
         self.assertTrue(self.collector.running)  # Still running
-        self.assertFalse(self.collector.paused)  # Not paused
+        self.assertTrue(self.collector.paused)  # Now paused
 
-        # Try pressing 'r' (reset) - should be ignored
-        old_total = self.collector.total_samples = 100
+        # Try pressing 'r' (reset) - should be ignored when finished
+        self.collector.total_samples = 100
         self.display.simulate_input(ord("r"))
         self.collector._handle_input()
-        self.assertEqual(self.collector.total_samples, old_total)  # Not reset
+        self.assertTrue(self.collector.running)  # Still running
+        self.assertEqual(self.collector.total_samples, 100)  # NOT reset when finished
 
         # Press 'q' - should stop
         self.display.simulate_input(ord("q"))
@@ -365,6 +368,35 @@ def test_finished_state_footer_message(self):
         # Check that footer contains finished message
         self.assertTrue(self.display.contains_text("PROFILING FINISHED"))
 
+    def test_finished_state_freezes_time(self):
+        """Test that time displays are frozen when finished."""
+        import time as time_module
+
+        # Set up collector with known start time
+        self.collector.start_time = time_module.perf_counter() - 10.0  # 10 seconds ago
+
+        # Mark as finished - this should freeze the time
+        self.collector.mark_finished()
+
+        # Get the frozen elapsed time
+        frozen_elapsed = self.collector.elapsed_time
+        frozen_time_display = self.collector.current_time_display
+
+        # Wait a bit to ensure time would advance
+        time_module.sleep(0.1)
+
+        # Time should remain frozen
+        self.assertEqual(self.collector.elapsed_time, frozen_elapsed)
+        self.assertEqual(self.collector.current_time_display, frozen_time_display)
+
+        # Verify finish timestamp was set
+        self.assertIsNotNone(self.collector.finish_timestamp)
+
+        # Reset should clear the frozen state
+        self.collector.reset_stats()
+        self.assertFalse(self.collector.finished)
+        self.assertIsNone(self.collector.finish_timestamp)
+
 
 class TestLiveCollectorFiltering(unittest.TestCase):
     """Tests for filtering functionality."""
@@ -1131,5 +1163,63 @@ def test_function_counts_are_per_thread_in_per_thread_mode(self):
         self.assertEqual(thread_222_funcs, {"func4", "func5"})
 
 
+class TestLiveCollectorNewFeatures(unittest.TestCase):
+    """Tests for new features added to live collector."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.display = MockDisplay()
+        self.collector = LiveStatsCollector(1000, display=self.display)
+        self.collector.start_time = time.perf_counter()
+
+    def test_filter_input_takes_precedence_over_commands(self):
+        """Test that filter input mode blocks command keys like 'h' and 'p'."""
+        # Enter filter input mode
+        self.collector.filter_input_mode = True
+        self.collector.filter_input_buffer = ""
+
+        # Press 'h' - should add to filter buffer, not show help
+        self.display.simulate_input(ord("h"))
+        self.collector._handle_input()
+
+        self.assertFalse(self.collector.show_help)  # Help not triggered
+        self.assertEqual(self.collector.filter_input_buffer, "h")  # Added to filter
+        self.assertTrue(self.collector.filter_input_mode)  # Still in filter mode
+
+    def test_reset_blocked_when_finished(self):
+        """Test that reset command is blocked when profiling is finished."""
+        # Set up some sample data and mark as finished
+        self.collector.total_samples = 100
+        self.collector.finished = True
+
+        # Press 'r' for reset
+        self.display.simulate_input(ord("r"))
+        self.collector._handle_input()
+
+        # Should NOT have been reset
+        self.assertEqual(self.collector.total_samples, 100)
+        self.assertTrue(self.collector.finished)
+
+    def test_time_display_fix_when_finished(self):
+        """Test that time display shows correct frozen time when finished."""
+        import time as time_module
+
+        # Mark as finished to freeze time
+        self.collector.mark_finished()
+
+        # Should have set both timestamps correctly
+        self.assertIsNotNone(self.collector.finish_timestamp)
+        self.assertIsNotNone(self.collector.finish_wall_time)
+
+        # Get the frozen time display
+        frozen_time = self.collector.current_time_display
+
+        # Wait a bit
+        time_module.sleep(0.1)
+
+        # Should still show the same frozen time (not jump to wrong time)
+        self.assertEqual(self.collector.current_time_display, frozen_time)
+
+
 if __name__ == "__main__":
     unittest.main()