fix: parse carriage-return progress updates in notebooks

Author: MilesCranmerBot

pysr/jupyter_progress.py

@@ -120,14 +120,26 @@ def __init__(self, target_stream, parser: _ProgressLineParser):
         self._parser = parser
         self._buffer = ""
 
+    def _drain_complete_lines(self) -> None:
+        # ProgressMeter often updates in-place with carriage returns (`\r`) rather
+        # than newline-terminated lines. Treat both as parse boundaries.
+        while True:
+            newline_idx = self._buffer.find("\n")
+            carriage_idx = self._buffer.find("\r")
+            candidates = [idx for idx in (newline_idx, carriage_idx) if idx != -1]
+            if not candidates:
+                break
+            split_idx = min(candidates)
+            line = self._buffer[:split_idx]
+            self._buffer = self._buffer[split_idx + 1 :]
+            self._parser.parse_line(line)
+
     def write(self, text: str) -> int:
         if not isinstance(text, str):
             text = str(text)
         written = self._target.write(text)
         self._buffer += text
-        while "\n" in self._buffer:
-            line, self._buffer = self._buffer.split("\n", 1)
-            self._parser.parse_line(line)
+        self._drain_complete_lines()
         return written if isinstance(written, int) else len(text)
 
     def flush(self) -> None:

pysr/test/test_jupyter_progress_helpers.py

@@ -48,6 +48,20 @@ def test_progress_parser_extracts_evolving_percent_lines(self):
         parser.parse_line("Evolving for 40 iterations... 11%|██        | ETA: 0:00:05")
         self.assertEqual(updates, [(4, 40)])
 
+    def test_progress_capture_stream_handles_carriage_return_updates(self):
+        module = _load_module()
+        updates = []
+        parser = module._ProgressLineParser(
+            lambda current, total: updates.append((current, total))
+        )
+        target = io.StringIO()
+        stream = module._ProgressCaptureStream(target, parser)
+        stream.write("Evolving for 100 iterations... 1%|\r")
+        stream.write("Evolving for 100 iterations... 2%|\r")
+        stream.write("Evolving for 100 iterations... 3%|\r")
+        stream.flush()
+        self.assertEqual(updates, [(1, 100), (2, 100), (3, 100)])
+
     def test_should_use_jupyter_progress_gating(self):
         module = _load_module()
         self.assertFalse(

scripts/test_jupyter_progress_helpers.py

@@ -45,10 +45,23 @@ def test_parser_extracts_evolving_percent() -> None:
     assert updates == [(4, 40)], updates
 
 
+def test_capture_stream_handles_carriage_return_updates() -> None:
+    module = _load_module()
+    updates = []
+    parser = module._ProgressLineParser(lambda current, total: updates.append((current, total)))
+    capture = module._ProgressCaptureStream(io.StringIO(), parser)
+    capture.write("Evolving for 100 iterations... 1%|\r")
+    capture.write("Evolving for 100 iterations... 2%|\r")
+    capture.write("Evolving for 100 iterations... 3%|\r")
+    capture.flush()
+    assert updates == [(1, 100), (2, 100), (3, 100)], updates
+
+
 def run() -> None:
     test_parser_extracts_progress()
     test_capture_stream_handles_split_lines()
     test_parser_extracts_evolving_percent()
+    test_capture_stream_handles_carriage_return_updates()
     print("helpers-tests=ok")