pythongh-118911: Trailing whitespace in a block shouldnt prevent the user from terminating the code block

Author: aelsayed95

Lib/_pyrepl/historical_reader.py

@@ -259,7 +259,7 @@ def select_item(self, i: int) -> None:
         self.transient_history[self.historyi] = self.get_unicode()
         buf = self.transient_history.get(i)
         if buf is None:
-            buf = self.history[i]
+            buf = self.history[i].rstrip()
         self.buffer = list(buf)
         self.historyi = i
         self.pos = len(self.buffer)

Lib/_pyrepl/readline.py

@@ -219,16 +219,29 @@ def do(self) -> None:
         # if there are already several lines and the cursor
         # is not on the last one, always insert a new \n.
         text = r.get_unicode()
+
         if "\n" in r.buffer[r.pos :] or (
             r.more_lines is not None and r.more_lines(text)
         ):
+            def _whitespace_before_pos():
+                before_idx = r.pos - 1
+                while before_idx > 0 and text[before_idx].isspace():
+                    before_idx -= 1
+                return text[before_idx : r.pos].count("\n") > 0
             #
-            # 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 there's already a new line before the cursor then
+            # even if the cursor is followed by whitespace, we assume
+            # the user is trying to terminate the block
+            if _whitespace_before_pos() and text[r.pos:].isspace():
+                self.finish = True
+            else:
+                #
+                # 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])
         elif not self.reader.paste_mode:
             self.finish = True
         else:

Lib/test/test_pyrepl/test_pyrepl.py

@@ -317,14 +317,13 @@ def test_multiline_edit(self):
             [
                 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="left", raw=bytearray(b"\x1bOD")),
+                Event(evt="key", data="left", raw=bytearray(b"\x1bOD")),
+                Event(evt="key", data="left", raw=bytearray(b"\x1bOD")),
                 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")),
-                Event(evt="key", data="down", raw=bytearray(b"\x1bOB")),
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
                 Event(evt="key", data="\n", raw=bytearray(b"\n")),
             ],
         )

Lib/test/test_pyrepl/test_reader.py

@@ -1,7 +1,8 @@
 import itertools
+import functools
 from unittest import TestCase
 
-from .support import handle_all_events, handle_events_narrow_console, code_to_events
+from .support import handle_all_events, handle_events_narrow_console, code_to_events, prepare_reader
 from _pyrepl.console import Event
 
 
@@ -133,3 +134,45 @@ def test_up_arrow_after_ctrl_r(self):
 
         reader, _ = handle_all_events(events)
         self.assert_screen_equals(reader, "")
+
+    def test_newline_within_block_trailing_whitespace(self):
+        # fmt: off
+        code = (
+            "def foo():\n"
+            "  a = 1\n"
+        )
+        # fmt: on
+
+        events = itertools.chain(
+            code_to_events(code),
+            [
+                # go to the end of the first line 
+                Event(evt="key", data="up", raw=bytearray(b"\x1bOA")),
+                Event(evt="key", data="up", raw=bytearray(b"\x1bOA")),
+                Event(evt="key", data="\x05", raw=bytearray(b"\x1bO5")),
+                # new lines in-block shouldn't terminate the block
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
+                # end of line 2
+                Event(evt="key", data="down", raw=bytearray(b"\x1bOB")),
+                Event(evt="key", data="\x05", raw=bytearray(b"\x1bO5")),
+                # a double new line in-block should terminate the block
+                # even if its followed by whitespace
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
+            ],
+        )
+
+        no_paste_reader = functools.partial(prepare_reader, paste_mode=False)
+        reader, _ = handle_all_events(events, prepare_reader=no_paste_reader)
+
+        expected = (
+            "def foo():\n"
+            "\n"
+            "\n"
+            "  a = 1\n"
+            "  \n"
+            "  "    # HistoricalReader will trim trailing whitespace
+        )
+        self.assert_screen_equals(reader, expected)
+        self.assertEqual(reader.finished, True)