gh-121610: pyrepl - handle extending blocks when multi-statement blocks are pasted (GH-121757)

saucoide · ambv · web-flow · commit 7d111dac160c · 2024-07-16T01:33:57.000+02:00
console.compile with the "single" param throws an exception when there are multiple statements, never allowing to adding newlines to a pasted code block (gh-121610) This add a few extra checks to allow extending when in an indented block, and tests for a few examples Co-authored-by: Łukasz Langa <lukasz@langa.pl>
diff --git a/Lib/_pyrepl/simple_interact.py b/Lib/_pyrepl/simple_interact.py
@@ -27,6 +27,7 @@
 
 import _sitebuiltins
 import linecache
+import functools
 import sys
 import code
 
@@ -78,6 +79,25 @@ def _clear_screen():
 }
 
 
+def _more_lines(console: code.InteractiveConsole, unicodetext: str) -> bool:
+    # ooh, look at the hack:
+    src = _strip_final_indent(unicodetext)
+    try:
+        code = console.compile(src, "<stdin>", "single")
+    except (OverflowError, SyntaxError, ValueError):
+        lines = src.splitlines(keepends=True)
+        if len(lines) == 1:
+            return False
+
+        last_line = lines[-1]
+        was_indented = last_line.startswith((" ", "\t"))
+        not_empty = last_line.strip() != ""
+        incomplete = not last_line.endswith("\n")
+        return (was_indented or not_empty) and incomplete
+    else:
+        return code is None
+
+
 def run_multiline_interactive_console(
     console: code.InteractiveConsole,
     *,
@@ -88,6 +108,7 @@ def run_multiline_interactive_console(
     if future_flags:
         console.compile.compiler.flags |= future_flags
 
+    more_lines = functools.partial(_more_lines, console)
     input_n = 0
 
     def maybe_run_command(statement: str) -> bool:
@@ -113,16 +134,6 @@ def maybe_run_command(statement: str) -> bool:
 
         return False
 
-    def more_lines(unicodetext: str) -> bool:
-        # ooh, look at the hack:
-        src = _strip_final_indent(unicodetext)
-        try:
-            code = console.compile(src, "<stdin>", "single")
-        except (OverflowError, SyntaxError, ValueError):
-            return False
-        else:
-            return code is None
-
     while 1:
         try:
             try:
diff --git a/Lib/test/test_pyrepl/test_interact.py b/Lib/test/test_pyrepl/test_interact.py
@@ -7,7 +7,7 @@
 from test.support import force_not_colorized
 
 from _pyrepl.console import InteractiveColoredConsole
-
+from _pyrepl.simple_interact import _more_lines
 
 class TestSimpleInteract(unittest.TestCase):
     def test_multiple_statements(self):
@@ -111,3 +111,104 @@ def test_no_active_future(self):
             result = console.runsource(source)
         self.assertFalse(result)
         self.assertEqual(f.getvalue(), "{'x': <class 'int'>}\n")
+
+
+class TestMoreLines(unittest.TestCase):
+    def test_invalid_syntax_single_line(self):
+        namespace = {}
+        code = "if foo"
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertFalse(_more_lines(console, code))
+
+    def test_empty_line(self):
+        namespace = {}
+        code = ""
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertFalse(_more_lines(console, code))
+
+    def test_valid_single_statement(self):
+        namespace = {}
+        code = "foo = 1"
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertFalse(_more_lines(console, code))
+
+    def test_multiline_single_assignment(self):
+        namespace = {}
+        code = dedent("""\
+        foo = [
+            1,
+            2,
+            3,
+        ]""")
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertFalse(_more_lines(console, code))
+
+    def test_multiline_single_block(self):
+        namespace = {}
+        code = dedent("""\
+        def foo():
+            '''docs'''
+
+            return 1""")
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertTrue(_more_lines(console, code))
+
+    def test_multiple_statements_single_line(self):
+        namespace = {}
+        code = "foo = 1;bar = 2"
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertFalse(_more_lines(console, code))
+
+    def test_multiple_statements(self):
+        namespace = {}
+        code = dedent("""\
+        import time
+
+        foo = 1""")
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertTrue(_more_lines(console, code))
+
+    def test_multiple_blocks(self):
+        namespace = {}
+        code = dedent("""\
+        from dataclasses import dataclass
+
+        @dataclass
+        class Point:
+            x: float
+            y: float""")
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertTrue(_more_lines(console, code))
+
+    def test_multiple_blocks_empty_newline(self):
+        namespace = {}
+        code = dedent("""\
+        from dataclasses import dataclass
+
+        @dataclass
+        class Point:
+            x: float
+            y: float
+        """)
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertFalse(_more_lines(console, code))
+
+    def test_multiple_blocks_indented_newline(self):
+        namespace = {}
+        code = (
+            "from dataclasses import dataclass\n"
+            "\n"
+            "@dataclass\n"
+            "class Point:\n"
+            "    x: float\n"
+            "    y: float\n"
+            "    "
+        )
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertFalse(_more_lines(console, code))
+
+    def test_incomplete_statement(self):
+        namespace = {}
+        code = "if foo:"
+        console = InteractiveColoredConsole(namespace, filename="<stdin>")
+        self.assertTrue(_more_lines(console, code))
