Separate _pyrepl overlays from base rendering

Move completion and message UI onto overlays layered over the base rendered screen.
The completion menu stops pretending it is the base reality.

Author: pablogsal

Lib/_pyrepl/completing_reader.py

@@ -24,7 +24,7 @@
 
 import re
 from . import commands, console, reader
-from .render import RenderLine, RenderedScreen
+from .render import RenderLine, ScreenOverlay
 from .reader import Reader
 
 
@@ -283,15 +283,15 @@ def after_command(self, cmd: Command) -> None:
         if not isinstance(cmd, (complete, self_insert)):
             self.cmpltn_reset()
 
-    def calc_screen(self) -> RenderedScreen:
-        rendered_screen = super().calc_screen()
-        if self.cmpltn_menu_visible:
-            rendered_screen = rendered_screen.with_overlay(
+    def get_screen_overlays(self) -> tuple[ScreenOverlay, ...]:
+        if not self.cmpltn_menu_visible:
+            return ()
+        return (
+            ScreenOverlay(
                 self.lxy[1] + 1,
-                (RenderLine.from_rendered_text(line) for line in self.cmpltn_menu),
-            )
-            self.rendered_screen = rendered_screen
-        return rendered_screen
+                tuple(RenderLine.from_rendered_text(line) for line in self.cmpltn_menu),
+            ),
+        )
 
     def finish(self) -> None:
         super().finish()

Lib/_pyrepl/reader.py

@@ -36,7 +36,7 @@
     process_prompt as build_prompt_content,
 )
 from .layout import LayoutMap, LayoutResult, LayoutRow, WrappedRow, layout_content_lines
-from .render import RenderCell, RenderLine, RenderedScreen, StyleRef
+from .render import RenderCell, RenderLine, RenderedScreen, ScreenOverlay, StyleRef
 from .utils import ANSI_ESCAPE_SEQUENCE, wlen, gen_colors
 from .trace import trace
 
@@ -380,7 +380,7 @@ def collect_keymap(self) -> tuple[tuple[KeySpec, CommandName], ...]:
         return default_keymap
 
     def calc_screen(self) -> RenderedScreen:
-        """Translate changes in self.buffer into a structured rendered screen."""
+        """Translate the editable buffer into a base rendered screen."""
         num_common_lines = 0
         offset = 0
         if self.last_refresh_cache.valid(self):
@@ -427,12 +427,7 @@ def calc_screen(self) -> RenderedScreen:
             layout_rows,
             last_refresh_line_end_offsets,
         )
-
-        render_lines = base_render_lines.copy()
-        render_lines.extend(self._render_message_lines())
-
-        self.rendered_screen = RenderedScreen(tuple(render_lines), self.cxy)
-        return self.rendered_screen
+        return RenderedScreen(tuple(base_render_lines), self.cxy)
 
     def _buffer_refresh_from_pos(self) -> int | None:
         buffer_from_pos = self.invalidation.buffer_rebuild_from_pos
@@ -536,9 +531,9 @@ def _render_wrapped_rows(
             for row in wrapped_rows
         ]
 
-    def _render_message_lines(self) -> list[RenderLine]:
+    def _render_message_lines(self) -> tuple[RenderLine, ...]:
         if not self.msg:
-            return []
+            return ()
         width = self.console.width
         render_lines: list[RenderLine] = []
         for message_line in self.msg.split("\n"):
@@ -551,7 +546,19 @@ def _render_message_lines(self) -> list[RenderLine]:
                 render_lines.append(
                     RenderLine.from_rendered_text(message_line[offset : offset + width])
                 )
-        return render_lines
+        return tuple(render_lines)
+
+    def get_screen_overlays(self) -> tuple[ScreenOverlay, ...]:
+        return ()
+
+    def compose_rendered_screen(self, base_screen: RenderedScreen) -> RenderedScreen:
+        overlays = list(self.get_screen_overlays())
+        message_lines = self._render_message_lines()
+        if message_lines:
+            overlays.append(ScreenOverlay(len(base_screen.lines), message_lines))
+        if not overlays:
+            return base_screen
+        return RenderedScreen(base_screen.lines, base_screen.cursor, tuple(overlays))
 
     def _render_line(
         self,
@@ -763,7 +770,8 @@ def prepare(self) -> None:
             self.rendered_screen = RenderedScreen.empty()
             self.invalidate_full()
             self.last_command = None
-            self.calc_screen()
+            base_screen = self.calc_screen()
+            self.rendered_screen = self.compose_rendered_screen(base_screen)
         except BaseException:
             self.restore()
             raise
@@ -821,7 +829,9 @@ def update_screen(self) -> None:
     def refresh(self) -> None:
         """Recalculate and refresh the screen."""
         # this call sets up self.cxy, so call it first.
-        rendered_screen = self.calc_screen()
+        base_screen = self.calc_screen()
+        rendered_screen = self.compose_rendered_screen(base_screen)
+        self.rendered_screen = rendered_screen
         trace(
             "reader.refresh cursor={cursor} lines={lines} "
             "dims=({width},{height}) invalidation={invalidation}",

Lib/test/test_pyrepl/test_reader.py

@@ -324,8 +324,41 @@ def test_message_refresh_keeps_base_render_cache(self):
         reader.update_screen()
 
         self.assertEqual([line.text for line in reader.last_refresh_cache.render_lines], ["ab"])
+        self.assertEqual([line.text for line in reader.rendered_screen.lines], ["ab"])
+        self.assertEqual(len(reader.rendered_screen.overlays), 1)
         self.assertEqual(reader.screen, ["ab", "! boom "])
 
+    def test_completion_overlay_keeps_base_render_cache(self):
+        namespace = {"itertools": itertools}
+        code = "itertools."
+        events = itertools.chain(
+            code_to_events(code),
+            [
+                Event(evt="key", data="\t", raw=bytearray(b"\t")),
+                Event(evt="key", data="\t", raw=bytearray(b"\t")),
+            ],
+        )
+
+        completing_reader = functools.partial(
+            prepare_reader,
+            readline_completer=rlcompleter.Completer(namespace).complete,
+        )
+        reader, _ = handle_all_events(events, prepare_reader=completing_reader)
+
+        self.assertEqual([line.text for line in reader.last_refresh_cache.render_lines], [code])
+        self.assertEqual([line.text for line in reader.rendered_screen.lines], [code])
+        self.assertEqual(len(reader.rendered_screen.overlays), 1)
+        self.assertEqual(reader.screen[0], code)
+        self.assertEqual(reader.screen[1].rstrip(), "itertools.accumulate(")
+
+        reader.cmpltn_reset()
+        reader.update_screen()
+
+        self.assertEqual([line.text for line in reader.last_refresh_cache.render_lines], [code])
+        self.assertEqual([line.text for line in reader.rendered_screen.lines], [code])
+        self.assertEqual(reader.rendered_screen.overlays, ())
+        self.assertEqual(reader.screen, [code])
+
     def test_completions_updated_on_key_press(self):
         namespace = {"itertools": itertools}
         code = "itertools."