Add support for if else statements

rewrite/rewrite/python/format/spaces_visitor.py

@@ -3,7 +3,7 @@
 import rewrite.java as j
 from rewrite import Tree
 from rewrite.java import J, Assignment, JLeftPadded, AssignmentOperation, MemberReference, MethodInvocation, \
-    MethodDeclaration, Empty, ArrayAccess, Space
+    MethodDeclaration, Empty, ArrayAccess, Space, If
 from rewrite.python import PythonVisitor, SpacesStyle, Binary, ChainedAssignment
 from rewrite.visitor import P
 
@@ -206,6 +206,28 @@ def _apply_binary_space_around(self, binary, use_space_around: bool):
         binary = binary.with_right(self.space_before(binary.right, use_space_around))
         return binary
 
+    def visit_if(self, if_stm: If, p: P) -> J:
+        if_: j.If = cast(If, super().visit_if(if_stm, p))
+
+        # Handle space before if condition e.g. if    True: <-> if True:
+        if_ = if_.with_if_condition(self.space_before(if_._if_condition, True))
+
+        # Handle space before if colon e.g. if True:    pass <-> if True: pass
+        if_ = if_.with_if_condition(
+            if_.if_condition.padding.with_tree(
+                SpacesVisitor.space_after(if_.if_condition.padding.tree, self._style.other.before_colon))
+        )
+
+        # TODO: Handle tree
+
+        return if_
+
+    def visit_else(self, else_: If.Else, p: P) -> J:
+        e: j.If.Else = cast(j.If.Else, super().visit_else(else_, p))
+        e = e.padding.with_body(self.space_before_right_padded_element(e.padding.body, False))
+
+        return e
+
     @staticmethod
     def space_before(j: J, space_before: bool) -> J:
         space: Space = cast(Space, j.prefix)
@@ -229,6 +251,10 @@ def space_before_container(container: j.JContainer, space_before: bool) -> j.JCo
         else:
             return container
 
+    @staticmethod
+    def space_before_right_padded_element(container: j.JRightPadded, space_before: bool) -> j.JRightPadded:
+        return container.with_element(SpacesVisitor.space_before(container.element, space_before))
+
     @staticmethod
     def space_last_comment_suffix(comments: List[j.Comment], space_suffix: bool) -> List[j.Comment]:
         return [SpacesVisitor.space_suffix(comment, space_suffix) if i == len(comments) - 1 else comment for i, comment

rewrite/tests/python/all/format/spaces_visitor_test.py

@@ -429,3 +429,88 @@ def test_spaces_python_binary_operators_string(left_space, right_space):
         spec=RecipeSpec()
         .with_recipe(from_visitor(SpacesVisitor(style)))
     )
+
+def test_if_statement():
+    style = IntelliJ.spaces()
+    rewrite_run(
+        # language=python
+        python(
+            """
+            if    1 >   2    :
+                pass
+            """,
+            """
+            if 1 > 2:
+                pass
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipe(from_visitor(SpacesVisitor(style)))
+    )
+
+def test_if_statement_multiple_conditions():
+    style = IntelliJ.spaces()
+    rewrite_run(
+        # language=python
+        python(
+            """
+            if    1 >   2  and   (2 < 3)  :
+                pass
+            """,
+            """
+            if 1 > 2 and (2 < 3):
+                pass
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipe(from_visitor(SpacesVisitor(style)))
+    )
+
+def test_if_else_statement():
+    style = IntelliJ.spaces()
+    rewrite_run(
+        # language=python
+        python(
+            """
+            if   1 > 2  :
+                a = 1
+            else  :
+                a = 2
+            """,
+            """
+            if 1 > 2:
+                a = 1
+            else:
+                a = 2
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipe(from_visitor(SpacesVisitor(style)))
+    )
+
+
+def test_if_elif_else_statement():
+    style = IntelliJ.spaces()
+    rewrite_run(
+        # language=python
+        python(
+            """
+            if   1 > 2  :
+                print(1)
+            elif    1 < 2  :
+                print(2)
+            else  :
+                print(3)
+            """,
+            """
+            if 1 > 2:
+                print(1)
+            elif 1 < 2:
+                print(2)
+            else:
+                print(3)
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipe(from_visitor(SpacesVisitor(style)))
+    )