requested changes

Author: tom-pytel

Lib/ast.py

@@ -708,10 +708,10 @@ def maybe_semicolon(self):
         if self._source:
             self.write("; ")
 
-    def fill(self, text="", allow_semi=True):
+    def fill(self, text="", *, allow_semicolon=True):
         """Indent a piece of text and append it, according to the current
         indentation level, or only delineate with semicolon if applicable"""
-        if self._in_interactive and not self._indent and allow_semi:
+        if self._in_interactive and not self._indent and allow_semicolon:
             self.maybe_semicolon()
             self.write(text)
         else:
@@ -960,17 +960,17 @@ def visit_Raise(self, node):
             self.traverse(node.cause)
 
     def do_visit_try(self, node):
-        self.fill("try", allow_semi=False)
+        self.fill("try", allow_semicolon=False)
         with self.block():
             self.traverse(node.body)
         for ex in node.handlers:
             self.traverse(ex)
         if node.orelse:
-            self.fill("else", allow_semi=False)
+            self.fill("else", allow_semicolon=False)
             with self.block():
                 self.traverse(node.orelse)
         if node.finalbody:
-            self.fill("finally", allow_semi=False)
+            self.fill("finally", allow_semicolon=False)
             with self.block():
                 self.traverse(node.finalbody)
 
@@ -991,7 +991,7 @@ def visit_TryStar(self, node):
             self._in_try_star = prev_in_try_star
 
     def visit_ExceptHandler(self, node):
-        self.fill("except*" if self._in_try_star else "except", allow_semi=False)
+        self.fill("except*" if self._in_try_star else "except", allow_semicolon=False)
         if node.type:
             self.write(" ")
             self.traverse(node.type)
@@ -1004,9 +1004,9 @@ def visit_ExceptHandler(self, node):
     def visit_ClassDef(self, node):
         self.maybe_newline()
         for deco in node.decorator_list:
-            self.fill("@", allow_semi=False)
+            self.fill("@", allow_semicolon=False)
             self.traverse(deco)
-        self.fill("class " + node.name, allow_semi=False)
+        self.fill("class " + node.name, allow_semicolon=False)
         if hasattr(node, "type_params"):
             self._type_params_helper(node.type_params)
         with self.delimit_if("(", ")", condition = node.bases or node.keywords):
@@ -1036,10 +1036,10 @@ def visit_AsyncFunctionDef(self, node):
     def _function_helper(self, node, fill_suffix):
         self.maybe_newline()
         for deco in node.decorator_list:
-            self.fill("@", allow_semi=False)
+            self.fill("@", allow_semicolon=False)
             self.traverse(deco)
         def_str = fill_suffix + " " + node.name
-        self.fill(def_str, allow_semi=False)
+        self.fill(def_str, allow_semicolon=False)
         if hasattr(node, "type_params"):
             self._type_params_helper(node.type_params)
         with self.delimit("(", ")"):
@@ -1090,54 +1090,54 @@ def visit_AsyncFor(self, node):
         self._for_helper("async for ", node)
 
     def _for_helper(self, fill, node):
-        self.fill(fill, allow_semi=False)
+        self.fill(fill, allow_semicolon=False)
         self.set_precedence(_Precedence.TUPLE, node.target)
         self.traverse(node.target)
         self.write(" in ")
         self.traverse(node.iter)
         with self.block(extra=self.get_type_comment(node)):
             self.traverse(node.body)
         if node.orelse:
-            self.fill("else", allow_semi=False)
+            self.fill("else", allow_semicolon=False)
             with self.block():
                 self.traverse(node.orelse)
 
     def visit_If(self, node):
-        self.fill("if ", allow_semi=False)
+        self.fill("if ", allow_semicolon=False)
         self.traverse(node.test)
         with self.block():
             self.traverse(node.body)
         # collapse nested ifs into equivalent elifs.
         while node.orelse and len(node.orelse) == 1 and isinstance(node.orelse[0], If):
             node = node.orelse[0]
-            self.fill("elif ", allow_semi=False)
+            self.fill("elif ", allow_semicolon=False)
             self.traverse(node.test)
             with self.block():
                 self.traverse(node.body)
         # final else
         if node.orelse:
-            self.fill("else", allow_semi=False)
+            self.fill("else", allow_semicolon=False)
             with self.block():
                 self.traverse(node.orelse)
 
     def visit_While(self, node):
-        self.fill("while ", allow_semi=False)
+        self.fill("while ", allow_semicolon=False)
         self.traverse(node.test)
         with self.block():
             self.traverse(node.body)
         if node.orelse:
-            self.fill("else", allow_semi=False)
+            self.fill("else", allow_semicolon=False)
             with self.block():
                 self.traverse(node.orelse)
 
     def visit_With(self, node):
-        self.fill("with ", allow_semi=False)
+        self.fill("with ", allow_semicolon=False)
         self.interleave(lambda: self.write(", "), self.traverse, node.items)
         with self.block(extra=self.get_type_comment(node)):
             self.traverse(node.body)
 
     def visit_AsyncWith(self, node):
-        self.fill("async with ", allow_semi=False)
+        self.fill("async with ", allow_semicolon=False)
         self.interleave(lambda: self.write(", "), self.traverse, node.items)
         with self.block(extra=self.get_type_comment(node)):
             self.traverse(node.body)
@@ -1279,7 +1279,7 @@ def visit_Name(self, node):
         self.write(node.id)
 
     def _write_docstring(self, node):
-        self.fill(allow_semi=False)
+        self.fill(allow_semicolon=False)
         if node.kind == "u":
             self.write("u")
         self._write_str_avoiding_backslashes(node.value, quote_types=_MULTI_QUOTES)
@@ -1573,7 +1573,7 @@ def visit_Slice(self, node):
             self.traverse(node.step)
 
     def visit_Match(self, node):
-        self.fill("match ", allow_semi=False)
+        self.fill("match ", allow_semicolon=False)
         self.traverse(node.subject)
         with self.block():
             for case in node.cases:
@@ -1667,7 +1667,7 @@ def visit_withitem(self, node):
             self.traverse(node.optional_vars)
 
     def visit_match_case(self, node):
-        self.fill("case ", allow_semi=False)
+        self.fill("case ", allow_semicolon=False)
         self.traverse(node.pattern)
         if node.guard:
             self.write(" if ")

Lib/test/test_ast/test_ast.py

@@ -814,17 +814,6 @@ def test_repr_large_input_crash(self):
                                     r"Exceeds the limit \(\d+ digits\)"):
             repr(ast.Constant(value=eval(source)))
 
-    def test_unparse_interactive(self):
-        # gh-129598: Fix of ast.unparse() when ast.Interactive contains multiple statements
-        source = "i = 1; 'expr'; raise Exception"
-        self.assertEqual(source, ast.unparse(ast.parse(source, mode='single')))
-        source = "if i:\n 'expr'\nelse:\n raise Exception"
-        unparsed = "if i:\n    'expr'\nelse:\n    raise Exception"
-        self.assertEqual(unparsed, ast.unparse(ast.parse(source, mode='single')))
-        source = "@decorator\ndef func():\n 'docstring'\n i = 1; 'expr'; raise Exception"
-        unparsed = '''@decorator\ndef func():\n    """docstring"""\n    i = 1\n    'expr'\n    raise Exception'''
-        self.assertEqual(unparsed, ast.unparse(ast.parse(source, mode='single')))
-
 
 class CopyTests(unittest.TestCase):
     """Test copying and pickling AST nodes."""

Lib/test/test_unparse.py

@@ -141,13 +141,13 @@ def check_invalid(self, node, raises=ValueError):
         with self.subTest(node=node):
             self.assertRaises(raises, ast.unparse, node)
 
-    def get_source(self, code1, code2=None):
+    def get_source(self, code1, code2=None, **kwargs):
         code2 = code2 or code1
-        code1 = ast.unparse(ast.parse(code1))
+        code1 = ast.unparse(ast.parse(code1, **kwargs))
         return code1, code2
 
-    def check_src_roundtrip(self, code1, code2=None):
-        code1, code2 = self.get_source(code1, code2)
+    def check_src_roundtrip(self, code1, code2=None, **kwargs):
+        code1, code2 = self.get_source(code1, code2, **kwargs)
         with self.subTest(code1=code1, code2=code2):
             self.assertEqual(code2, code1)
 
@@ -466,6 +466,112 @@ def test_type_ignore(self):
         ):
             self.check_ast_roundtrip(statement, type_comments=True)
 
+    def test_unparse_interactive(self):
+        # gh-129598: Fix of ast.unparse() when ast.Interactive contains multiple statements
+        self.check_src_roundtrip("i = 1; 'expr'; raise Exception", mode='single')
+        self.check_src_roundtrip("i: int = 1; j: float = 0; k += l", mode='single')
+        combinable = (
+            "'expr'",
+            "(i := 1)",
+            "import foo",
+            "from foo import bar",
+            "i = 1",
+            "i += 1",
+            "i: int = 1",
+            "return i",
+            "pass",
+            "break",
+            "continue",
+            "del i",
+            "assert i",
+            "global i",
+            "nonlocal j",
+            "await i",
+            "yield i",
+            "yield from i",
+            "raise i",
+            "type t[T] = ...",
+            "i",
+        )
+        for a in combinable:
+            for b in combinable:
+                self.check_src_roundtrip(f"{a}; {b}", mode='single')
+
+        # rest of the tests just make sure mode='single' parse and unparse didn't break
+        self.check_src_roundtrip(
+            "if i:\n 'expr'\nelse:\n raise Exception",
+            "if i:\n    'expr'\nelse:\n    raise Exception",
+            mode='single'
+        )
+        self.check_src_roundtrip(
+            "@decorator1\n@decorator2\ndef func():\n 'docstring'\n i = 1; 'expr'; raise Exception",
+            '''@decorator1\n@decorator2\ndef func():\n    """docstring"""\n    i = 1\n    'expr'\n    raise Exception''',
+            mode='single'
+        )
+        self.check_src_roundtrip(
+            "@decorator1\n@decorator2\nclass cls:\n 'docstring'\n i = 1; 'expr'; raise Exception",
+            '''@decorator1\n@decorator2\nclass cls:\n    """docstring"""\n    i = 1\n    'expr'\n    raise Exception''',
+            mode='single'
+        )
+        for statement in (
+            "def x():\n    pass",
+            "def x(y):\n    pass",
+            "async def x():\n    pass",
+            "async def x(y):\n    pass",
+            "for x in y:\n    pass",
+            "async for x in y:\n    pass",
+            "with x():\n    pass",
+            "async with x():\n    pass",
+            "def f():\n    pass",
+            "def f(a):\n    pass",
+            "def f(b=2):\n    pass",
+            "def f(a, b):\n    pass",
+            "def f(a, b=2):\n    pass",
+            "def f(a=5, b=2):\n    pass",
+            "def f(*, a=1, b=2):\n    pass",
+            "def f(*, a=1, b):\n    pass",
+            "def f(*, a, b=2):\n    pass",
+            "def f(a, b=None, *, c, **kwds):\n    pass",
+            "def f(a=2, *args, c=5, d, **kwds):\n    pass",
+            "def f(*args, **kwargs):\n    pass",
+            "class cls:\n\n    def f(self):\n        pass",
+            "class cls:\n\n    def f(self, a):\n        pass",
+            "class cls:\n\n    def f(self, b=2):\n        pass",
+            "class cls:\n\n    def f(self, a, b):\n        pass",
+            "class cls:\n\n    def f(self, a, b=2):\n        pass",
+            "class cls:\n\n    def f(self, a=5, b=2):\n        pass",
+            "class cls:\n\n    def f(self, *, a=1, b=2):\n        pass",
+            "class cls:\n\n    def f(self, *, a=1, b):\n        pass",
+            "class cls:\n\n    def f(self, *, a, b=2):\n        pass",
+            "class cls:\n\n    def f(self, a, b=None, *, c, **kwds):\n        pass",
+            "class cls:\n\n    def f(self, a=2, *args, c=5, d, **kwds):\n        pass",
+            "class cls:\n\n    def f(self, *args, **kwargs):\n        pass",
+        ):
+            self.check_src_roundtrip(statement, mode='single')
+        for statement in (
+            "def x():",
+            "def x(y):",
+            "async def x():",
+            "async def x(y):",
+            "for x in y:",
+            "async for x in y:",
+            "with x():",
+            "async with x():",
+            "def f():",
+            "def f(a):",
+            "def f(b=2):",
+            "def f(a, b):",
+            "def f(a, b=2):",
+            "def f(a=5, b=2):",
+            "def f(*, a=1, b=2):",
+            "def f(*, a=1, b):",
+            "def f(*, a, b=2):",
+            "def f(a, b=None, *, c, **kwds):",
+            "def f(a=2, *args, c=5, d, **kwds):",
+            "def f(*args, **kwargs):",
+        ):
+            self.check_src_roundtrip(statement + '\n i=1;j=2', statement + '\n    i = 1\n    j = 2', mode='single')
+
 
 class CosmeticTestCase(ASTTestCase):
     """Test if there are cosmetic issues caused by unnecessary additions"""