gh-131507: Add support for syntax highlighting in PyREPL (GH-133247)

Co-authored-by: Victorien <65306057+Viicos@users.noreply.github.com>
Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>

5 files changed, 221 insertions, 19 deletions

diff --git a/Lib/test/test_pyrepl/test_pyrepl.py b/Lib/test/test_pyrepl/test_pyrepl.py
index 75a5afa..93029ab 100644
--- a/Lib/test/test_pyrepl/test_pyrepl.py
+++ b/Lib/test/test_pyrepl/test_pyrepl.py
@@ -45,6 +45,7 @@ class ReplTestCase(TestCase):
         cmdline_args: list[str] | None = None,
         cwd: str | None = None,
         skip: bool = False,
+        timeout: float = SHORT_TIMEOUT,
     ) -> tuple[str, int]:
         temp_dir = None
         if cwd is None:
@@ -52,7 +53,12 @@ class ReplTestCase(TestCase):
             cwd = temp_dir.name
         try:
             return self._run_repl(
-                repl_input, env=env, cmdline_args=cmdline_args, cwd=cwd, skip=skip,
+                repl_input,
+                env=env,
+                cmdline_args=cmdline_args,
+                cwd=cwd,
+                skip=skip,
+                timeout=timeout,
             )
         finally:
             if temp_dir is not None:
@@ -66,6 +72,7 @@ class ReplTestCase(TestCase):
         cmdline_args: list[str] | None,
         cwd: str,
         skip: bool,
+        timeout: float,
     ) -> tuple[str, int]:
         assert pty
         master_fd, slave_fd = pty.openpty()
@@ -103,7 +110,7 @@ class ReplTestCase(TestCase):
         os.write(master_fd, repl_input.encode("utf-8"))
 
         output = []
-        while select.select([master_fd], [], [], SHORT_TIMEOUT)[0]:
+        while select.select([master_fd], [], [], timeout)[0]:
             try:
                 data = os.read(master_fd, 1024).decode("utf-8")
                 if not data:
@@ -114,12 +121,12 @@ class ReplTestCase(TestCase):
         else:
             os.close(master_fd)
             process.kill()
-            process.wait(timeout=SHORT_TIMEOUT)
+            process.wait(timeout=timeout)
             self.fail(f"Timeout while waiting for output, got: {''.join(output)}")
 
         os.close(master_fd)
         try:
-            exit_code = process.wait(timeout=SHORT_TIMEOUT)
+            exit_code = process.wait(timeout=timeout)
         except subprocess.TimeoutExpired:
             process.kill()
             exit_code = process.wait()
@@ -1561,25 +1568,29 @@ class TestMain(ReplTestCase):
 
     def test_history_survive_crash(self):
         env = os.environ.copy()
-        commands = "1\nexit()\n"
-        output, exit_code = self.run_repl(commands, env=env, skip=True)
 
         with tempfile.NamedTemporaryFile() as hfile:
             env["PYTHON_HISTORY"] = hfile.name
-            commands = "spam\nimport time\ntime.sleep(1000)\npreved\n"
+
+            commands = "1\n2\n3\nexit()\n"
+            output, exit_code = self.run_repl(commands, env=env, skip=True)
+
+            commands = "spam\nimport time\ntime.sleep(1000)\nquit\n"
             try:
-                self.run_repl(commands, env=env)
+                self.run_repl(commands, env=env, timeout=3)
             except AssertionError:
                 pass
 
             history = pathlib.Path(hfile.name).read_text()
+            self.assertIn("2", history)
+            self.assertIn("exit()", history)
             self.assertIn("spam", history)
-            self.assertIn("time", history)
+            self.assertIn("import time", history)
             self.assertNotIn("sleep", history)
-            self.assertNotIn("preved", history)
+            self.assertNotIn("quit", history)
 
     def test_keyboard_interrupt_after_isearch(self):
-        output, exit_code = self.run_repl(["\x12", "\x03", "exit"])
+        output, exit_code = self.run_repl("\x12\x03exit\n")
         self.assertEqual(exit_code, 0)
 
     def test_prompt_after_help(self):
diff --git a/Lib/test/test_pyrepl/test_reader.py b/Lib/test/test_pyrepl/test_reader.py
index 109cb60..8d7fcf5 100644
--- a/Lib/test/test_pyrepl/test_reader.py
+++ b/Lib/test/test_pyrepl/test_reader.py
@@ -1,14 +1,21 @@
 import itertools
 import functools
 import rlcompleter
+from textwrap import dedent
 from unittest import TestCase
 from unittest.mock import MagicMock
 
 from .support import handle_all_events, handle_events_narrow_console
 from .support import ScreenEqualMixin, code_to_events
-from .support import prepare_reader, prepare_console
+from .support import prepare_console, reader_force_colors
+from .support import reader_no_colors as prepare_reader
 from _pyrepl.console import Event
 from _pyrepl.reader import Reader
+from _colorize import theme
+
+
+overrides = {"RESET": "z", "SOFT_KEYWORD": "K"}
+colors = {overrides.get(k, k[0].lower()): v for k, v in theme.items()}
 
 
 class TestReader(ScreenEqualMixin, TestCase):
@@ -123,8 +130,9 @@ class TestReader(ScreenEqualMixin, TestCase):
     def test_control_characters(self):
         code = 'flag = "🏳️‍🌈"'
         events = code_to_events(code)
-        reader, _ = handle_all_events(events)
+        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
         self.assert_screen_equal(reader, 'flag = "🏳️\\u200d🌈"', clean=True)
+        self.assert_screen_equal(reader, 'flag {o}={z} {s}"🏳️\\u200d🌈"{z}'.format(**colors))
 
     def test_setpos_from_xy_multiple_lines(self):
         # fmt: off
@@ -355,3 +363,140 @@ class TestReader(ScreenEqualMixin, TestCase):
         reader, _ = handle_all_events(events)
         reader.setpos_from_xy(8, 0)
         self.assertEqual(reader.pos, 7)
+
+    def test_syntax_highlighting_basic(self):
+        code = dedent(
+            """\
+            import re, sys
+            def funct(case: str = sys.platform) -> None:
+                match = re.search(
+                    "(me)",
+                    '''
+                    Come on
+                      Come on now
+                        You know that it's time to emerge
+                    ''',
+                )
+                match case:
+                    case "emscripten": print("on the web")
+                    case "ios" | "android": print("on the phone")
+                    case _: print('arms around', match.group(1))
+            """
+        )
+        expected = dedent(
+            """\
+            {k}import{z} re{o},{z} sys
+            {a}{k}def{z} {d}funct{z}{o}({z}case{o}:{z} {b}str{z} {o}={z} sys{o}.{z}platform{o}){z} {o}->{z} {k}None{z}{o}:{z}
+                match {o}={z} re{o}.{z}search{o}({z}
+                    {s}"(me)"{z}{o},{z}
+                    {s}'''{z}
+            {s}        Come on{z}
+            {s}          Come on now{z}
+            {s}            You know that it's time to emerge{z}
+            {s}        '''{z}{o},{z}
+                {o}){z}
+                {K}match{z} case{o}:{z}
+                    {K}case{z} {s}"emscripten"{z}{o}:{z} {b}print{z}{o}({z}{s}"on the web"{z}{o}){z}
+                    {K}case{z} {s}"ios"{z} {o}|{z} {s}"android"{z}{o}:{z} {b}print{z}{o}({z}{s}"on the phone"{z}{o}){z}
+                    {K}case{z} {K}_{z}{o}:{z} {b}print{z}{o}({z}{s}'arms around'{z}{o},{z} match{o}.{z}group{o}({z}{n}1{z}{o}){z}{o}){z}
+            """
+        )
+        expected_sync = expected.format(a="", **colors)
+        events = code_to_events(code)
+        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        self.assert_screen_equal(reader, code, clean=True)
+        self.assert_screen_equal(reader, expected_sync)
+        self.assertEqual(reader.pos, 2**7 + 2**8)
+        self.assertEqual(reader.cxy, (0, 14))
+
+        async_msg = "{k}async{z} ".format(**colors)
+        expected_async = expected.format(a=async_msg, **colors)
+        more_events = itertools.chain(
+            code_to_events(code),
+            [Event(evt="key", data="up", raw=bytearray(b"\x1bOA"))] * 13,
+            code_to_events("async "),
+        )
+        reader, _ = handle_all_events(more_events, prepare_reader=reader_force_colors)
+        self.assert_screen_equal(reader, expected_async)
+        self.assertEqual(reader.pos, 21)
+        self.assertEqual(reader.cxy, (6, 1))
+
+    def test_syntax_highlighting_incomplete_string_first_line(self):
+        code = dedent(
+            """\
+            def unfinished_function(arg: str = "still typing
+            """
+        )
+        expected = dedent(
+            """\
+            {k}def{z} {d}unfinished_function{z}{o}({z}arg{o}:{z} {b}str{z} {o}={z} {s}"still typing{z}
+            """
+        ).format(**colors)
+        events = code_to_events(code)
+        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        self.assert_screen_equal(reader, code, clean=True)
+        self.assert_screen_equal(reader, expected)
+
+    def test_syntax_highlighting_incomplete_string_another_line(self):
+        code = dedent(
+            """\
+            def unfinished_function(
+                arg: str = "still typing
+            """
+        )
+        expected = dedent(
+            """\
+            {k}def{z} {d}unfinished_function{z}{o}({z}
+                arg{o}:{z} {b}str{z} {o}={z} {s}"still typing{z}
+            """
+        ).format(**colors)
+        events = code_to_events(code)
+        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        self.assert_screen_equal(reader, code, clean=True)
+        self.assert_screen_equal(reader, expected)
+
+    def test_syntax_highlighting_incomplete_multiline_string(self):
+        code = dedent(
+            """\
+            def unfinished_function():
+                '''Still writing
+                the docstring
+            """
+        )
+        expected = dedent(
+            """\
+            {k}def{z} {d}unfinished_function{z}{o}({z}{o}){z}{o}:{z}
+                {s}'''Still writing{z}
+            {s}    the docstring{z}
+            """
+        ).format(**colors)
+        events = code_to_events(code)
+        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        self.assert_screen_equal(reader, code, clean=True)
+        self.assert_screen_equal(reader, expected)
+
+    def test_syntax_highlighting_incomplete_fstring(self):
+        code = dedent(
+            """\
+            def unfinished_function():
+                var = f"Single-quote but {
+                1
+                +
+                1
+                } multi-line!
+            """
+        )
+        expected = dedent(
+            """\
+            {k}def{z} {d}unfinished_function{z}{o}({z}{o}){z}{o}:{z}
+                var {o}={z} {s}f"{z}{s}Single-quote but {z}{o}{OB}{z}
+                {n}1{z}
+                {o}+{z}
+                {n}1{z}
+                {o}{CB}{z}{s} multi-line!{z}
+            """
+        ).format(OB="{", CB="}", **colors)
+        events = code_to_events(code)
+        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        self.assert_screen_equal(reader, code, clean=True)
+        self.assert_screen_equal(reader, expected)
diff --git a/Lib/test/test_pyrepl/test_unix_console.py b/Lib/test/test_pyrepl/test_unix_console.py
index 2f5c150..7acb84a 100644
--- a/Lib/test/test_pyrepl/test_unix_console.py
+++ b/Lib/test/test_pyrepl/test_unix_console.py
@@ -33,10 +33,12 @@ def unix_console(events, **kwargs):
 
 handle_events_unix_console = partial(
     handle_all_events,
-    prepare_console=partial(unix_console),
+    prepare_reader=reader_no_colors,
+    prepare_console=unix_console,
 )
 handle_events_narrow_unix_console = partial(
     handle_all_events,
+    prepare_reader=reader_no_colors,
     prepare_console=partial(unix_console, width=5),
 )
 handle_events_short_unix_console = partial(
diff --git a/Lib/test/test_pyrepl/test_utils.py b/Lib/test/test_pyrepl/test_utils.py
index 0d59968..8ce1e53 100644
--- a/Lib/test/test_pyrepl/test_utils.py
+++ b/Lib/test/test_pyrepl/test_utils.py
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from _pyrepl.utils import str_width, wlen
+from _pyrepl.utils import str_width, wlen, prev_next_window
 
 
 class TestUtils(TestCase):
@@ -25,3 +25,38 @@ class TestUtils(TestCase):
 
         self.assertEqual(wlen('hello'), 5)
         self.assertEqual(wlen('hello' + '\x1a'), 7)
+
+    def test_prev_next_window(self):
+        def gen_normal():
+            yield 1
+            yield 2
+            yield 3
+            yield 4
+
+        pnw = prev_next_window(gen_normal())
+        self.assertEqual(next(pnw), (None, 1, 2))
+        self.assertEqual(next(pnw), (1, 2, 3))
+        self.assertEqual(next(pnw), (2, 3, 4))
+        self.assertEqual(next(pnw), (3, 4, None))
+        with self.assertRaises(StopIteration):
+            next(pnw)
+
+        def gen_short():
+            yield 1
+
+        pnw = prev_next_window(gen_short())
+        self.assertEqual(next(pnw), (None, 1, None))
+        with self.assertRaises(StopIteration):
+            next(pnw)
+
+        def gen_raise():
+            yield from gen_normal()
+            1/0
+
+        pnw = prev_next_window(gen_raise())
+        self.assertEqual(next(pnw), (None, 1, 2))
+        self.assertEqual(next(pnw), (1, 2, 3))
+        self.assertEqual(next(pnw), (2, 3, 4))
+        self.assertEqual(next(pnw), (3, 4, None))
+        with self.assertRaises(ZeroDivisionError):
+            next(pnw)
diff --git a/Lib/test/test_pyrepl/test_windows_console.py b/Lib/test/test_pyrepl/test_windows_console.py
index 69f2d5a..e95fec4 100644
--- a/Lib/test/test_pyrepl/test_windows_console.py
+++ b/Lib/test/test_pyrepl/test_windows_console.py
@@ -12,6 +12,7 @@ from unittest import TestCase
 from unittest.mock import MagicMock, call
 
 from .support import handle_all_events, code_to_events
+from .support import reader_no_colors as default_prepare_reader
 
 try:
     from _pyrepl.console import Event, Console
@@ -47,14 +48,22 @@ class WindowsConsoleTests(TestCase):
             setattr(console, key, val)
         return console
 
-    def handle_events(self, events: Iterable[Event], **kwargs):
-        return handle_all_events(events, partial(self.console, **kwargs))
+    def handle_events(
+        self,
+        events: Iterable[Event],
+        prepare_console=None,
+        prepare_reader=None,
+        **kwargs,
+    ):
+        prepare_console = prepare_console or partial(self.console, **kwargs)
+        prepare_reader = prepare_reader or default_prepare_reader
+        return handle_all_events(events, prepare_console, prepare_reader)
 
     def handle_events_narrow(self, events):
         return self.handle_events(events, width=5)
 
-    def handle_events_short(self, events):
-        return self.handle_events(events, height=1)
+    def handle_events_short(self, events, **kwargs):
+        return self.handle_events(events, height=1, **kwargs)
 
     def handle_events_height_3(self, events):
         return self.handle_events(events, height=3)