gh-118911: Trailing whitespace in a block shouldn't prevent the user from terminating the code block (#119355)

Co-authored-by: Łukasz Langa <lukasz@langa.pl>

5 files changed, 79 insertions, 9 deletions

diff --git a/Lib/_pyrepl/historical_reader.py b/Lib/_pyrepl/historical_reader.py
index eef7d90..121de33 100644
--- a/Lib/_pyrepl/historical_reader.py
+++ b/Lib/_pyrepl/historical_reader.py
@@ -259,7 +259,7 @@ class HistoricalReader(Reader):
         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)
diff --git a/Lib/_pyrepl/readline.py b/Lib/_pyrepl/readline.py
index 796f1ef..ffa14a9 100644
--- a/Lib/_pyrepl/readline.py
+++ b/Lib/_pyrepl/readline.py
@@ -244,14 +244,27 @@ class maybe_accept(commands.Command):
         r: ReadlineAlikeReader
         r = self.reader  # type: ignore[assignment]
         r.dirty = True  # this is needed to hide the completion menu, if visible
-        #
+
         # 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 _newline_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
+
+            # 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 _newline_before_pos() and text[r.pos:].isspace():
+                self.finish = True
+                return
+
             # auto-indent the next line like the previous line
             prevlinestart, indent = _get_previous_line_indent(r.buffer, r.pos)
             r.insert("\n")
diff --git a/Lib/test/test_pyrepl/test_pyrepl.py b/Lib/test/test_pyrepl/test_pyrepl.py
index 7b5217e..bdcabf9 100644
--- a/Lib/test/test_pyrepl/test_pyrepl.py
+++ b/Lib/test/test_pyrepl/test_pyrepl.py
@@ -405,12 +405,21 @@ class TestPyReplOutput(TestCase):
             [
                 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="backspace", raw=bytearray(b"\x7f")),
+                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"\x08")),
                 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="backspace", raw=bytearray(b"\x08")),
+                Event(evt="key", data="delete", raw=bytearray(b"\x7F")),
+                Event(evt="key", data="right", raw=bytearray(b"g")),
+                Event(evt="key", data="backspace", raw=bytearray(b"\x08")),
+                Event(evt="key", data="p", raw=bytearray(b"p")),
+                Event(evt="key", data="a", raw=bytearray(b"a")),
+                Event(evt="key", data="s", raw=bytearray(b"s")),
+                Event(evt="key", data="s", raw=bytearray(b"s")),
+                Event(evt="key", data="\n", raw=bytearray(b"\n")),
                 Event(evt="key", data="\n", raw=bytearray(b"\n")),
             ],
         )
@@ -419,7 +428,7 @@ class TestPyReplOutput(TestCase):
         output = multiline_input(reader)
         self.assertEqual(output, "def f():\n    ...\n    ")
         output = multiline_input(reader)
-        self.assertEqual(output, "def g():\n    ...\n    ")
+        self.assertEqual(output, "def g():\n    pass\n    ")
 
     def test_history_navigation_with_up_arrow(self):
         events = itertools.chain(
diff --git a/Lib/test/test_pyrepl/test_reader.py b/Lib/test/test_pyrepl/test_reader.py
index dc7d8a5..7bf7a36 100644
--- a/Lib/test/test_pyrepl/test_reader.py
+++ b/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 @@ class TestReader(TestCase):
 
         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.assertTrue(reader.finished)
diff --git a/Misc/NEWS.d/next/Library/2024-05-21-20-13-23.gh-issue-118911.iG8nMq.rst b/Misc/NEWS.d/next/Library/2024-05-21-20-13-23.gh-issue-118911.iG8nMq.rst
new file mode 100644
index 0000000..4f15c1b
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2024-05-21-20-13-23.gh-issue-118911.iG8nMq.rst
@@ -0,0 +1,5 @@
+In PyREPL, updated ``maybe-accept``'s logic so that if the user hits
+:kbd:`Enter` twice, they are able to terminate the block even if there's
+trailing whitespace. Also, now when the user hits arrow up, the cursor
+is on the last functional line. This matches IPython's behavior.
+Patch by Aya Elsayed.