gh-111201: Improve pyrepl auto indentation

Lib/_pyrepl/readline.py

@@ -230,13 +230,24 @@ def _get_first_indentation(buffer: list[str]) -> str | None:
     return None
 
 
-def _is_last_char_colon(buffer: list[str]) -> bool:
-    i = len(buffer)
-    while i > 0:
-        i -= 1
-        if buffer[i] not in " \t\n":  # ignore whitespaces
-            return buffer[i] == ":"
-    return False
+def _should_auto_indent(buffer: list[str], pos: int) -> bool:
+    # check if last character before "pos" is a colon, ignoring
+    # whitespaces and comments.
+    last_char = None
+    while pos > 0:
+        pos -= 1
+        if last_char is None:
+            if buffer[pos] not in " \t\n":  # ignore whitespaces
+                last_char = buffer[pos]
+        else:
+            # even if we found a non-whitespace character before
+            # original pos, we keep going back until newline is reached
+            # to make sure we ignore comments
+            if buffer[pos] == "\n":
+                break
+            if buffer[pos] == "#":
+                last_char = None
+    return last_char == ":"
 
 
 class maybe_accept(commands.Command):
@@ -273,7 +284,7 @@ def _newline_before_pos():
                     for i in range(prevlinestart, prevlinestart + indent):
                         r.insert(r.buffer[i])
                 r.update_last_used_indentation()
-                if _is_last_char_colon(r.buffer):
+                if _should_auto_indent(r.buffer, r.pos):
                     if r.last_used_indentation is not None:
                         indentation = r.last_used_indentation
                     else:

Lib/test/test_pyrepl/test_pyrepl.py

@@ -312,6 +312,14 @@ def test_cursor_position_after_wrap_and_move_up(self):
         self.assertEqual(reader.pos, 10)
         self.assertEqual(reader.cxy, (1, 1))
 
+
+class TestPyReplAutoindent(TestCase):
+    def prepare_reader(self, events):
+        console = FakeConsole(events)
+        config = ReadlineConfig(readline_completer=None)
+        reader = ReadlineAlikeReader(console=console, config=config)
+        return reader
+
     def test_auto_indent_default(self):
         # fmt: off
         input_code = (
@@ -372,7 +380,6 @@ def test_auto_indent_prev_block(self):
             ),
         )
 
-
         output_code = (
             "def g():\n"
             "  pass\n"
@@ -385,6 +392,78 @@ def test_auto_indent_prev_block(self):
         output2 = multiline_input(reader)
         self.assertEqual(output2, output_code)
 
+    def test_auto_indent_multiline(self):
+        # fmt: off
+        events = itertools.chain(
+            code_to_events(
+                "def f():\n"
+                    "pass"
+            ),
+            [
+                # go to the end of the first line
+                Event(evt="key", data="up", raw=bytearray(b"\x1bOA")),
+                Event(evt="key", data="\x05", raw=bytearray(b"\x1bO5")),
+                # new line should be autoindented
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
+            ],
+            code_to_events(
+                "pass"
+            ),
+            [
+                # go to end of last line
+                Event(evt="key", data="down", raw=bytearray(b"\x1bOB")),
+                Event(evt="key", data="\x05", raw=bytearray(b"\x1bO5")),
+                # double newline to terminate the block
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
+            ],
+        )
+
+        output_code = (
+            "def f():\n"
+            "    pass\n"
+            "    pass\n"
+            "    "
+        )
+        # fmt: on
+
+        reader = self.prepare_reader(events)
+        output = multiline_input(reader)
+        self.assertEqual(output, output_code)
+
+    def test_auto_indent_with_comment(self):
+        # fmt: off
+        events = code_to_events(
+            "def f():  # foo\n"
+                "pass\n\n"
+        )
+
+        output_code = (
+            "def f():  # foo\n"
+            "    pass\n"
+            "    "
+        )
+        # fmt: on
+
+        reader = self.prepare_reader(events)
+        output = multiline_input(reader)
+        self.assertEqual(output, output_code)
+
+    def test_auto_indent_ignore_comments(self):
+        # fmt: off
+        events = code_to_events(
+            "pass  #:\n"
+        )
+
+        output_code = (
+            "pass  #:"
+        )
+        # fmt: on
+
+        reader = self.prepare_reader(events)
+        output = multiline_input(reader)
+        self.assertEqual(output, output_code)
+
 
 class TestPyReplOutput(TestCase):
     def prepare_reader(self, events):

Lib/test/test_pyrepl/test_reader.py

@@ -168,8 +168,8 @@ def test_newline_within_block_trailing_whitespace(self):
 
         expected = (
             "def foo():\n"
-            "\n"
-            "\n"
+            "    \n"
+            "    \n"
             "    a = 1\n"
             "    \n"
             "    "    # HistoricalReader will trim trailing whitespace