Autoformat/spacevisitor

rewrite/rewrite/markers.py

@@ -3,12 +3,12 @@
 import traceback
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, replace
-from typing import List, ClassVar, cast, TYPE_CHECKING, Optional
+from typing import List, ClassVar, cast, TYPE_CHECKING, Callable, TypeVar, Type, Optional
 from uuid import UUID
 
 if TYPE_CHECKING:
     from .parser import Parser
-from .utils import random_id
+from .utils import random_id, list_map
 
 
 class Marker(ABC):
@@ -30,6 +30,8 @@ def __hash__(self) -> int:
         return hash(self.id)
 
 
+M = TypeVar('M', bound=Marker)
+
 @dataclass(frozen=True, eq=False)
 class Markers:
     _id: UUID
@@ -50,14 +52,24 @@ def markers(self) -> List[Marker]:
     def with_markers(self, markers: List[Marker]) -> Markers:
         return self if markers is self._markers else Markers(self._id, markers)
 
-    def find_first(self, type: type) -> Optional[Marker]:
+    def find_first(self, cls: Type[M]) -> Optional[M]:
         for marker in self.markers:
-            if isinstance(marker, type):
+            if isinstance(marker, cls):
                 return marker
         return None
 
-    def find_all(self, type: type) -> List[Marker]:
-        return [m for m in self.markers if isinstance(m, type)]
+    def find_all(self, cls: Type[M]) -> List[M]:
+        return [m for m in self.markers if isinstance(m, cls)]
+
+    def compute_by_type(self, cls: Type[M], remap_fn: Callable[[M], Marker]) -> Markers:
+        """
+        Replace all markers of the given type with the result of the function.
+
+        :param cls: type of the markers to remap
+        :param remap_fn: function to remap the marker
+        :return: new Markers instance with the updated markers, or the same instance if no markers were updated
+        """
+        return self.with_markers(list_map(lambda m: remap_fn(m) if isinstance(m, cls) else m, self.markers))
 
     EMPTY: ClassVar[Markers]
 

rewrite/rewrite/python/format/auto_format.py

@@ -1,13 +1,14 @@
-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 .normalize_tabs_or_spaces import NormalizeTabsOrSpacesVisitor
 from .. import TabsAndIndentsStyle
 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
 
 
@@ -32,29 +33,3 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) ->
             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)