auto-indentation in _pyrepl

Author: wiggin15

Lib/_pyrepl/readline.py

@@ -98,6 +98,7 @@ class ReadlineAlikeReader(historical_reader.HistoricalReader, CompletingReader):
     # Instance fields
     config: ReadlineConfig
     more_lines: MoreLinesCallable | None = None
+    last_used_indentation: str | None = None
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -156,6 +157,11 @@ def get_trimmed_history(self, maxlength: int) -> list[str]:
             cut = 0
         return self.history[cut:]
 
+    def update_last_used_indentation(self) -> None:
+        indentation = _get_first_indentation(self.buffer)
+        if indentation is not None:
+            self.last_used_indentation = indentation
+
     # --- simplified support for reading multiline Python statements ---
 
     def collect_keymap(self) -> tuple[tuple[KeySpec, CommandName], ...]:
@@ -210,6 +216,28 @@ def _get_previous_line_indent(buffer: list[str], pos: int) -> tuple[int, int | N
     return prevlinestart, indent
 
 
+def _get_first_indentation(buffer: list[str]) -> str | None:
+    indented_line_start = None
+    for i in range(len(buffer)):
+        if (i < len(buffer) - 1
+            and buffer[i] == "\n"
+            and buffer[i + 1] in " \t"
+        ):
+            indented_line_start = i + 1
+        elif indented_line_start is not None and buffer[i] not in " \t\n":
+            return ''.join(buffer[indented_line_start : i])
+    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
+
+
 class maybe_accept(commands.Command):
     def do(self) -> None:
         r: ReadlineAlikeReader
@@ -226,9 +254,18 @@ def do(self) -> None:
             # auto-indent the next line like the previous line
             prevlinestart, indent = _get_previous_line_indent(r.buffer, r.pos)
             r.insert("\n")
-            if not self.reader.paste_mode and indent:
-                for i in range(prevlinestart, prevlinestart + indent):
-                    r.insert(r.buffer[i])
+            if not self.reader.paste_mode:
+                if indent:
+                    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 r.last_used_indentation is not None:
+                        indentation = r.last_used_indentation
+                    else:
+                        # default
+                        indentation = " " * 4
+                    r.insert(indentation)
         elif not self.reader.paste_mode:
             self.finish = True
         else:

Lib/test/test_pyrepl/test_pyrepl.py

@@ -4,7 +4,13 @@
 import unittest
 from unittest import TestCase
 
-from .support import FakeConsole, handle_all_events, handle_events_narrow_console, multiline_input, code_to_events
+from .support import (
+    FakeConsole,
+    handle_all_events,
+    handle_events_narrow_console,
+    multiline_input,
+    code_to_events,
+)
 from _pyrepl.console import Event
 from _pyrepl.readline import ReadlineAlikeReader, ReadlineConfig
 
@@ -313,14 +319,12 @@ def test_basic(self):
 
     def test_multiline_edit(self):
         events = itertools.chain(
-            code_to_events("def f():\n  ...\n\n"),
+            code_to_events("def f():\n...\n\n"),
             [
                 Event(evt="key", data="up", raw=bytearray(b"\x1bOA")),
                 Event(evt="key", data="up", raw=bytearray(b"\x1bOA")),
                 Event(evt="key", data="up", raw=bytearray(b"\x1bOA")),
                 Event(evt="key", data="right", raw=bytearray(b"\x1bOC")),
-                Event(evt="key", data="right", raw=bytearray(b"\x1bOC")),
-                Event(evt="key", data="right", raw=bytearray(b"\x1bOC")),
                 Event(evt="key", data="backspace", raw=bytearray(b"\x7f")),
                 Event(evt="key", data="g", raw=bytearray(b"g")),
                 Event(evt="key", data="down", raw=bytearray(b"\x1bOB")),
@@ -331,9 +335,9 @@ def test_multiline_edit(self):
         reader = self.prepare_reader(events)
 
         output = multiline_input(reader)
-        self.assertEqual(output, "def f():\n  ...\n  ")
+        self.assertEqual(output, "def f():\n    ...\n    ")
         output = multiline_input(reader)
-        self.assertEqual(output, "def g():\n  ...\n  ")
+        self.assertEqual(output, "def g():\n    ...\n    ")
 
     def test_history_navigation_with_up_arrow(self):
         events = itertools.chain(
@@ -537,14 +541,14 @@ def test_paste_mid_newlines_not_in_paste_mode(self):
         # fmt: off
         code = (
             'def f():\n'
-            '  x = y\n'
-            '  \n'
-            '  y = z\n\n'
+                'x = y\n'
+                '\n'
+                'y = z\n\n'
         )
 
         expected = (
             'def f():\n'
-            '  x = y\n'
+            '    x = y\n'
             '    '
         )
         # fmt: on
@@ -558,19 +562,19 @@ def test_paste_not_in_paste_mode(self):
         # fmt: off
         input_code = (
             'def a():\n'
-            '  for x in range(10):\n'
-            '    if x%2:\n'
-            '      print(x)\n'
-            '    else:\n'
-            '      pass\n\n'
+                'for x in range(10):\n'
+                    'if x%2:\n'
+                        'print(x)\n'
+                    'else:\n'
+                        'pass\n\n'
         )
 
         output_code = (
             'def a():\n'
-            '  for x in range(10):\n'
-            '      if x%2:\n'
+            '    for x in range(10):\n'
+            '        if x%2:\n'
             '            print(x)\n'
-            '                else:'
+            '            else:'
         )
         # fmt: on
 
@@ -634,6 +638,81 @@ def test_bracketed_paste_single_line(self):
         output = multiline_input(reader)
         self.assertEqual(output, input_code)
 
+    def test_auto_indent_default(self):
+        # fmt: off
+        input_code = (
+            'def f():\n'
+                'pass\n\n'
+        )
+
+        output_code = (
+            'def f():\n'
+            '    pass\n'
+            '    '
+        )
+        # fmt: on
+
+    def test_auto_indent_continuation(self):
+        # auto indenting according to previous user indentation
+        # fmt: off
+        events = itertools.chain(
+            code_to_events("def f():\n"),
+            # add backspace to delete default auto-indent
+            [
+                Event(evt="key", data="backspace", raw=bytearray(b"\x7f")),
+            ],
+            code_to_events(
+                "  pass\n"
+                  "pass\n\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_prev_block(self):
+        # auto indenting according to indentation in different block
+        # fmt: off
+        events = itertools.chain(
+            code_to_events("def f():\n"),
+            # add backspace to delete default auto-indent
+            [
+                Event(evt="key", data="backspace", raw=bytearray(b"\x7f")),
+            ],
+            code_to_events(
+                "  pass\n"
+                "pass\n\n"
+            ),
+            code_to_events(
+                'def g():\n'
+                  'pass\n\n'
+            ),
+        )
+
+
+        output_code = (
+            'def g():\n'
+            '  pass\n'
+            '  '
+        )
+
+        # fmt: on
+
+        reader = self.prepare_reader(events)
+        output1 = multiline_input(reader)
+        output2 = multiline_input(reader)
+        self.assertEqual(output2, output_code)
+
 
 if __name__ == "__main__":
     unittest.main()