Autoformat/spacevisitor

rewrite/rewrite/python/format/auto_format.py

@@ -1,11 +1,12 @@
-from typing import Optional, cast, TypeVar
+from typing import Optional
 
 from .blank_lines import BlankLinesVisitor
 from .normalize_format import NormalizeFormatVisitor
+from .spaces_visitor import SpacesVisitor
 from ..style import BlankLinesStyle, SpacesStyle, IntelliJ
 from ..visitor import PythonVisitor
 from ... import Recipe, Tree, Cursor
-from ...java import JavaSourceFile, MethodDeclaration, J, Space
+from ...java import JavaSourceFile
 from ...visitor import P, T
 
 
@@ -23,32 +24,9 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) ->
         cu = tree if isinstance(tree, JavaSourceFile) else self._cursor.first_enclosing_or_throw(JavaSourceFile)
 
         tree = NormalizeFormatVisitor(self._stop_after).visit(tree, p, self._cursor.fork())
-        tree = BlankLinesVisitor(cu.get_style(BlankLinesStyle) or IntelliJ.blank_lines(), self._stop_after).visit(tree, p, self._cursor.fork())
-        tree = SpacesVisitor(cu.get_style(SpacesStyle) or IntelliJ.spaces(), self._stop_after).visit(tree, p, self._cursor.fork())
+        tree = BlankLinesVisitor(cu.get_style(BlankLinesStyle) or IntelliJ.blank_lines(), self._stop_after).visit(tree,
+                                                                                                                  p,
+                                                                                                                  self._cursor.fork())
+        tree = SpacesVisitor(cu.get_style(SpacesStyle) or IntelliJ.spaces(), self._stop_after).visit(tree, p,
+                                                                                                     self._cursor.fork())
         return tree
-
-
-J2 = TypeVar('J2', bound=J)
-
-
-class SpacesVisitor(PythonVisitor):
-    def __init__(self, style: SpacesStyle, stop_after: Tree = None):
-        self._style = style
-        self._before_parentheses = style.before_parentheses
-        self._stop_after = stop_after
-
-    def visit_method_declaration(self, md: MethodDeclaration, p: P) -> J:
-        md: MethodDeclaration = cast(MethodDeclaration, super().visit_method_declaration(md, p))
-        return md.padding.with_parameters(
-            md.padding.parameters.with_before(
-                Space.SINGLE_SPACE if self._before_parentheses.method_declaration else Space.EMPTY
-            )
-        )
-
-    def space_before(self, j: J2, space_before: bool) -> J2:
-        space: Space = cast(Space, j.prefix)
-        if space.comments or '\\' in space.whitespace:
-            # don't touch whitespaces with comments or continuation characters
-            return j
-
-        return j.with_prefix(Space.SINGLE_SPACE if space_before else Space.EMPTY)

rewrite/rewrite/python/format/spaces_visitor.py

@@ -0,0 +1,166 @@
+from typing import Optional, cast, TypeVar
+
+from rewrite import Tree
+from rewrite.java import J, Assignment, JLeftPadded, AssignmentOperation
+from rewrite.java import MethodInvocation, MethodDeclaration, Empty, ArrayAccess, Space
+from rewrite.python import PythonVisitor, SpacesStyle
+from rewrite.visitor import P
+
+J2 = TypeVar('J2', bound=J)
+
+
+class SpacesVisitor(PythonVisitor):
+    def __init__(self, style: SpacesStyle, stop_after: Tree = None):
+        self._style = style
+        self._before_parentheses = style.before_parentheses
+        self._stop_after = stop_after
+
+    def visit_method_declaration(self, method_declaration: MethodDeclaration, p: P) -> J:
+        md: MethodDeclaration = cast(MethodDeclaration, super().visit_method_declaration(method_declaration, p))
+        return md.padding.with_parameters(
+            md.padding.parameters.with_before(
+                Space.SINGLE_SPACE if self._before_parentheses.method_declaration else Space.EMPTY
+            )
+        )
+
+    def visit_method_invocation(self, method_invocation: MethodInvocation, p: P) -> J:
+        m: MethodInvocation = cast(MethodInvocation, super().visit_method_invocation(method_invocation, p))
+
+        # Handle space before parenthesis for method e.g. foo (..) <-> foo(..)
+        m = m.padding.with_arguments(m.padding.arguments.with_before(
+            Space.SINGLE_SPACE if self._style.before_parentheses.method_call else Space.EMPTY))
+
+        if not m.arguments or isinstance(m.arguments[0], Empty):
+            # Handle within method call parenthesis with no arguments e.g. foo() <-> foo( )
+            use_space = self._style.within.empty_method_call_parentheses
+            m = m.padding.with_arguments(
+                m.padding.arguments.padding.with_elements(
+                    [arg.with_element(self.space_before(arg.element, use_space)) for arg in
+                     m.padding.arguments.padding.elements]
+                )
+            )
+        else:
+            # Handle:
+            # - Within method call parenthesis with arguments e.g. foo( 1, 2 ) <-> foo(1, 2).
+            #   Note: this only applies to the space left of the first argument and right of the last argument.
+            # - Space after comma e.g. foo(1,2,3) <-> foo(1, 2, 3)
+            args_size = len(m.arguments)
+            use_space = self._style.within.method_call_parentheses
+
+            def _process_argument(index, arg, args_size, use_space):
+                if index == 0:
+                    arg = arg.with_element(self.space_before(arg.element, use_space))
+                else:
+                    arg = arg.with_element(self.space_before(arg.element, self._style.other.after_comma))
+
+                if index == args_size - 1:
+                    arg = self.space_after(arg, use_space)
+                else:
+                    arg = self.space_after(arg, self._style.other.before_comma)
+
+                return arg
+
+            m = m.padding.with_arguments(
+                m.padding.arguments.padding.with_elements(
+                    [_process_argument(index, arg, args_size, use_space)
+                     for index, arg in enumerate(m.padding.arguments.padding.elements)]
+                )
+            )
+        # TODO: Handle type parameters, relevant for constructors in Python.
+        return m
+
+    def visit_array_access(self, array_access: ArrayAccess, p: P) -> J:
+        a: ArrayAccess = cast(ArrayAccess, super().visit_array_access(array_access, p))
+        use_space_within_brackets = self._style.within.brackets
+        index_padding = a.dimension.padding.index
+        element_prefix = update_space(index_padding.element.prefix, use_space_within_brackets)
+        index_after = update_space(index_padding.after, use_space_within_brackets)
+
+        a = a.with_dimension(
+            a.dimension.padding.with_index(
+                index_padding.with_element(
+                    index_padding.element.with_prefix(element_prefix)
+                ).with_after(index_after)
+            ).with_prefix(update_space(a.dimension.prefix, self._style.before_parentheses.left_bracket))
+        )
+        return a
+
+    def visit_assignment(self, assignment: Assignment, p: P) -> J:
+        """
+        Handle assignment operator e.g. a = 1 <-> a=1
+        """
+        a: Assignment = cast(Assignment, super().visit_assignment(assignment, p))
+        a = a.padding.with_assignment(
+            self.space_before_jleftpadded(a.padding.assignment, self._style.around_operators.assignment))
+        a = a.padding.with_assignment(
+            a.padding.assignment.with_element(
+                self.space_before(a.padding.assignment.element, self._style.around_operators.assignment)))
+        return a
+
+    def visit_assignment_operation(self, assignment_operation: AssignmentOperation, p: P) -> J:
+        """
+        Handle assignment operation e.g. a += 1 <-> a+=1
+        """
+        a: AssignmentOperation = cast(AssignmentOperation, super().visit_assignment_operation(assignment_operation, p))
+        operator: JLeftPadded = a.padding.operator
+        a = a.padding.with_operator(
+            operator.with_before(update_space(operator.before, self._style.around_operators.assignment)))
+        return a.with_assignment(self.space_before(a.assignment, self._style.around_operators.assignment))
+
+    def space_before(self, j: J2, space_before: bool) -> J2:
+        space: Space = cast(Space, j.prefix)
+        if space.comments or '\\' in space.whitespace:
+            # don't touch whitespaces with comments or continuation characters
+            return j
+
+        return j.with_prefix(Space.SINGLE_SPACE if space_before else Space.EMPTY)
+
+    def space_before_jleftpadded(self, j: JLeftPadded[J2], space_before) -> JLeftPadded[J2]:
+        space: Space = cast(Space, j.before)
+        if space.comments or '\\' in space.whitespace:
+            # don't touch whitespaces with comments or continuation characters
+            return j
+
+        if space_before and not_single_space(space.whitespace):
+            return j.with_before(space.with_whitespace(" "))
+        elif not space_before and only_spaces_and_not_empty(space.whitespace):
+            return j.with_before(space.with_whitespace(""))
+        return j
+
+    def contains_comments(self, space: Space) -> bool:
+        return space.comments is not None or '\\' in space.whitespace
+
+    def space_after(self, j: J2, space_after: bool) -> J2:
+        space: Space = cast(Space, j.after)
+        if space.comments or '\\' in space.whitespace:
+            # don't touch whitespaces with comments or continuation characters
+            return j
+
+        if space_after and not_single_space(space.whitespace):
+            return j.with_after(space.with_whitespace(" "))
+        elif not space_after and only_spaces_and_not_empty(space.whitespace):
+            return j.with_after(space.with_whitespace(""))
+        return j
+
+def update_space(s: Space, have_space: bool) -> Space:
+    if s.comments:
+        return s
+
+    if have_space and not_single_space(s.whitespace):
+        return s.with_whitespace(" ")
+    elif not have_space and only_spaces_and_not_empty(s.whitespace):
+        return s.with_whitespace("")
+    else:
+        return s
+
+
+def only_spaces(s: Optional[str]) -> bool:
+    return s is not None and all(c in {' ', '\t'} for c in s)
+
+
+def only_spaces_and_not_empty(s: Optional[str]) -> bool:
+    return bool(s) and only_spaces(s)
+
+
+def not_single_space(s: Optional[str]) -> bool:
+    return s is not None and only_spaces(s) and s != " "

rewrite/tests/manual_parser_test.py

@@ -19,7 +19,7 @@
 RemotePrinterFactory(remoting_context.client).set_current()
 remoting_context.client.reset()
 
-parser = PythonParser()
+parser = PythonParser(None)
 res = parser.parse_inputs([ParserInput(pathlib.Path(path), None, False, lambda: open(path, 'r', newline='', encoding='utf-8'))], relative_to=None, ctx=InMemoryExecutionContext())
 for r in res:
     if isinstance(r, ParseError):