Correct semicolon terminated if-then statements

openrewrite · Nov 19, 2024 · ec1ba4f · ec1ba4f
1 parent 8d9a5a0
commit ec1ba4f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/rewrite/rewrite/python/_parser_visitor.py b/rewrite/rewrite/python/_parser_visitor.py
@@ -425,9 +425,8 @@ def visit_If(self, node):
         condition = j.ControlParentheses(random_id(), self.__whitespace(), Markers.EMPTY,
                                          self.__pad_right(self.__convert(node.test), self.__source_before(
                                              ':') if single_statement_then_body else Space.EMPTY))
-        then = self.__pad_right(
-            self.__convert(node.body[0]) if single_statement_then_body else self.__convert_block(node.body),
-            Space.EMPTY)
+        then = self.__pad_statement(node.body[0]) if single_statement_then_body else self.__pad_right(
+            self.__convert_block(node.body), Space.EMPTY)
         elze = None
         if len(node.orelse) > 0:
             else_prefix = self.__whitespace()
@@ -443,9 +442,8 @@ def visit_If(self, node):
                 random_id(),
                 else_prefix,
                 Markers.EMPTY,
-                self.__pad_right(
-                    self.__convert(node.orelse[0]) if single_statement_else_body else self.__convert_block(node.orelse),
-                    Space.EMPTY
+                self.__pad_statement(node.orelse[0]) if single_statement_else_body else self.__pad_right(
+                    self.__convert_block(node.orelse), Space.EMPTY
                 )
             )
         return j.If(

diff --git a/rewrite/tests/python/all/if_test.py b/rewrite/tests/python/all/if_test.py
@@ -15,6 +15,17 @@ def test_if():
     )
 
 
+def test_semicolon_terminated_then():
+    # language=python
+    rewrite_run(
+        python(r"""
+            if True:  i = 1;
+            else:     i = 2;
+            """
+       )
+    )
+
+
 def test_else_single():
     # language=python
     rewrite_run(