Add CoordinateBuilder.Block.last_statement() coordinate (#138)

This allows adding a statement as the last statement to a block.

Adding this required adding a new `MinimumViableSpacingVisitor` to add a linebreak prefix to any statements that don't already have one and where the preceding statement (if any) doesn't have a terminating semicolon.

Author: knutwannheden

rewrite/rewrite/__init__.py

@@ -7,7 +7,7 @@
 from .execution import ExecutionContext, DelegatingExecutionContext, InMemoryExecutionContext, Recipe, RecipeRunException
 from .markers import *
 from .tree import Checksum, FileAttributes, SourceFile, Tree, PrintOutputCapture, PrinterFactory
-from .utils import random_id, list_map, list_map_last
+from .utils import random_id, list_find, list_map, list_map_last
 from .visitor import Cursor, TreeVisitor
 from .parser import *
 from .result import *
@@ -21,6 +21,7 @@
     'PrintOutputCapture',
     'PrinterFactory',
     'random_id',
+    'list_find',
     'list_map',
     'list_map_last',
     'Cursor',

rewrite/rewrite/java/support_types.py

@@ -357,8 +357,15 @@ def replace(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
         return JavaCoordinates(self.tree, loc or Space.Location.STATEMENT_PREFIX, JavaCoordinates.Mode.REPLACE)
 
 
+@dataclass
+class _BlockCoordinateBuilder(_StatementCoordinateBuilder):
+    def last_statement(self) -> JavaCoordinates:
+        return self.before(Space.Location.BLOCK_END)
+
+
 CoordinateBuilder.Expression = _ExpressionCoordinateBuilder  # type: ignore
 CoordinateBuilder.Statement = _StatementCoordinateBuilder  # type: ignore
+CoordinateBuilder.Block = _BlockCoordinateBuilder  # type: ignore
 
 
 @dataclass

rewrite/rewrite/java/tree.py

@@ -733,6 +733,9 @@ def padding(self) -> PaddingHelper:
     def accept_java(self, v: JavaVisitor[P], p: P) -> J:
         return v.visit_block(self, p)
 
+    def get_coordinates(self) -> CoordinateBuilder.Block:
+        return CoordinateBuilder.Block(self)
+
 # noinspection PyShadowingBuiltins,PyShadowingNames,DuplicatedCode
 @dataclass(frozen=True, eq=False)
 class Break(Statement):
@@ -1289,8 +1292,10 @@ def padding(self) -> PaddingHelper:
         return p
 
     def printer(self, cursor: Cursor) -> TreeVisitor[Tree, PrintOutputCapture[P]]:
-        factory = PrinterFactory.current()
-        return factory.create_printer(cursor) if factory else JavaPrinter[PrintOutputCapture[P]]()
+        if factory := PrinterFactory.current():
+            return factory.create_printer(cursor)
+        from .printer import JavaPrinter
+        return JavaPrinter[PrintOutputCapture[P]]()
 
     def accept_java(self, v: JavaVisitor[P], p: P) -> J:
         return v.visit_compilation_unit(self, p)

rewrite/rewrite/python/__init__.py

@@ -70,6 +70,7 @@
     # Formatter
     'AutoFormat',
     'BlankLinesVisitor',
+    'MinimumViableSpacingVisitor',
     'NormalizeFormatVisitor',
     'NormalizeTabsOrSpacesVisitor',
     'SpacesVisitor',

rewrite/rewrite/python/_parser_visitor.py

@@ -422,28 +422,25 @@ def visit_While(self, node):
 
     def visit_If(self, node):
         prefix = self.__source_before('if')
-        single_statement_then_body = len(node.body) == 1
         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_statement(node.body[0]) if single_statement_then_body else self.__pad_right(
-            self.__convert_block(node.body), Space.EMPTY)
+                                         self.__pad_right(self.__convert(node.test), Space.EMPTY))
+        then = self.__pad_right(self.__convert_block(node.body), Space.EMPTY)
         elze = None
         if len(node.orelse) > 0:
             else_prefix = self.__whitespace()
             if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If) and self._source.startswith('elif',
                                                                                                         self._cursor):
-                single_statement_else_body = True
+                is_elif = True
                 self._cursor += 2
             else:
-                single_statement_else_body = False
+                is_elif = False
                 self._cursor += 4
 
             elze = j.If.Else(
                 random_id(),
                 else_prefix,
                 Markers.EMPTY,
-                self.__pad_statement(node.orelse[0]) if single_statement_else_body else self.__pad_right(
+                self.__pad_statement(node.orelse[0]) if is_elif else self.__pad_right(
                     self.__convert_block(node.orelse), Space.EMPTY
                 )
             )
@@ -2067,7 +2064,7 @@ def __next_lexer_token(self, tokens: Iterator[TokenInfo]) -> TokenInfo:
     def __convert_all(self, trees: Sequence) -> List[J2]:
         return [self.__convert(tree) for tree in trees]
 
-    def __convert_block(self, statements: Sequence, prefix: str = ':') -> j.Block:
+    def __convert_block(self, statements: Sequence[Statement], prefix: str = ':') -> j.Block:
         prefix = self.__source_before(prefix)
         if statements:
             statements = [self.__pad_statement(cast(ast.stmt, s)) for s in statements]

rewrite/rewrite/python/format/__init__.py

@@ -7,6 +7,7 @@
 __all__ = [
     'AutoFormat',
     'BlankLinesVisitor',
+    'MinimumViableSpacingVisitor',
     'NormalizeFormatVisitor',
     'NormalizeTabsOrSpacesVisitor',
     'SpacesVisitor',

rewrite/rewrite/python/format/auto_format.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 from .blank_lines import BlankLinesVisitor
+from .minimum_viable_spacing import MinimumViableSpacingVisitor
 from .normalize_format import NormalizeFormatVisitor
 from .normalize_line_breaks_visitor import NormalizeLineBreaksVisitor
 from .normalize_tabs_or_spaces import NormalizeTabsOrSpacesVisitor
@@ -31,6 +32,8 @@ def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) ->
 
         tree = NormalizeFormatVisitor(self._stop_after).visit(tree, p, self._cursor.fork())
 
+        tree = MinimumViableSpacingVisitor(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())

rewrite/rewrite/python/format/minimum_viable_spacing.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import cast, Optional, TypeVar, List
+
+from rewrite import Tree, P, Cursor, list_map, list_find
+from rewrite.java import J, Space, Statement, Block, Semicolon
+from rewrite.python import PythonVisitor, PyComment
+from rewrite.visitor import T
+
+J2 = TypeVar('J2', bound=J)
+
+
+class MinimumViableSpacingVisitor(PythonVisitor):
+    def __init__(self, stop_after: Optional[Tree] = None):
+        self._stop_after = stop_after
+        self._stop = False
+
+    def post_visit(self, tree: T, p: P) -> Optional[T]:
+        if self._stop_after and tree == self._stop_after:
+            self._stop = True
+
+        owner = self.cursor.parent_tree_cursor().value
+        if isinstance(tree, Statement) and isinstance(owner, Block) and not tree.prefix.comments and not '\n' in tree.prefix.whitespace:
+            statement_index = list_find(owner.statements, tree)
+            previous_statement = owner.padding.statements[statement_index - 1] if statement_index > 0 else None
+            if not previous_statement or not previous_statement.markers.find_first(Semicolon):
+                tree = tree.with_prefix(tree.prefix.with_whitespace('\n' + tree.prefix.whitespace))
+
+        return tree
+
+    def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> Optional[T]:
+        return tree if self._stop else super().visit(tree, p, parent)
+
+
+def _common_margin(s1, s2):
+    if s1 is None:
+        s = str(s2)
+        return s[s.rfind('\n') + 1:]
+
+    min_length = min(len(s1), len(s2))
+    for i in range(min_length):
+        if s1[i] != s2[i] or not s1[i].isspace():
+            return s1[:i]
+
+    return s2 if len(s2) < len(s1) else s1
+
+
+def _concatenate_prefix(j: J, prefix: Space) -> J2:
+    shift = _common_margin(None, j.prefix.whitespace)
+
+    def modify_comment(c: PyComment) -> PyComment:
+        if len(shift) == 0:
+            return c
+        c = c.with_text(c.text.replace('\n', '\n' + shift))
+        if '\n' in c.suffix:
+            c = c.with_suffix(c.suffix.replace('\n', '\n' + shift))
+        return c
+
+    comments = j.prefix.comments + list_map(modify_comment, cast(List[PyComment], prefix.comments))
+
+    new_prefix = j.prefix
+    new_prefix = new_prefix.with_whitespace(new_prefix.whitespace + prefix.whitespace)
+    if comments:
+        new_prefix = new_prefix.with_comments(comments)
+
+    return j.with_prefix(new_prefix)

rewrite/rewrite/python/templating.py

@@ -150,6 +150,11 @@ def visit_expression(self, expression: Expression, p: int) -> J:
         return expression
 
     def visit_block(self, block: Block, p: P) -> J:
+        if self.loc == Space.Location.BLOCK_END and block.is_scope(self.insertion_point):
+            parsed = self.template_parser.parse_block_statements(Cursor(self.cursor, self.insertion_point), Statement,
+                                                                 self.substituted_template, self.loc, self.mode)
+            gen: List[Statement] = self.substitutions.unsubstitute_all(parsed)
+            return self.auto_format(block.with_statements(block.statements + gen), p, self.cursor.parent) if gen else block
         if self.loc == Space.Location.STATEMENT_PREFIX:
             return self.auto_format(block.with_statements(list_flat_map(lambda s: self.get_replacements(s) if s.is_scope(self.insertion_point) else s, block.statements)), p, self.cursor.parent)
         return super().visit_block(block, p)

rewrite/rewrite/python/tree.py

@@ -96,13 +96,13 @@ def expression(self) -> Expression:
     def with_expression(self, expression: Expression) -> Await:
         return self if expression is self._expression else replace(self, _expression=expression)
 
-    _type: JavaType
+    _type: Optional[JavaType]
 
     @property
-    def type(self) -> JavaType:
+    def type(self) -> Optional[JavaType]:
         return self._type
 
-    def with_type(self, type: JavaType) -> Await:
+    def with_type(self, type: Optional[JavaType]) -> Await:
         return self if type is self._type else replace(self, _type=type)
 
     def accept_python(self, v: PythonVisitor[P], p: P) -> J:
@@ -340,13 +340,13 @@ def markers(self) -> Markers:
     def with_markers(self, markers: Markers) -> ExceptionType:
         return self if markers is self._markers else replace(self, _markers=markers)
 
-    _type: JavaType
+    _type: Optional[JavaType]
 
     @property
-    def type(self) -> JavaType:
+    def type(self) -> Optional[JavaType]:
         return self._type
 
-    def with_type(self, type: JavaType) -> ExceptionType:
+    def with_type(self, type: Optional[JavaType]) -> ExceptionType:
         return self if type is self._type else replace(self, _type=type)
 
     _exception_group: bool
@@ -503,13 +503,13 @@ def literal(self) -> Expression:
     def with_literal(self, literal: Expression) -> LiteralType:
         return self if literal is self._literal else replace(self, _literal=literal)
 
-    _type: JavaType
+    _type: Optional[JavaType]
 
     @property
-    def type(self) -> JavaType:
+    def type(self) -> Optional[JavaType]:
         return self._type
 
-    def with_type(self, type: JavaType) -> LiteralType:
+    def with_type(self, type: Optional[JavaType]) -> LiteralType:
         return self if type is self._type else replace(self, _type=type)
 
     def accept_python(self, v: PythonVisitor[P], p: P) -> J:
@@ -554,13 +554,13 @@ def type_tree(self) -> Expression:
     def with_type_tree(self, type_tree: Expression) -> TypeHint:
         return self if type_tree is self._type_tree else replace(self, _type_tree=type_tree)
 
-    _type: JavaType
+    _type: Optional[JavaType]
 
     @property
-    def type(self) -> JavaType:
+    def type(self) -> Optional[JavaType]:
         return self._type
 
-    def with_type(self, type: JavaType) -> TypeHint:
+    def with_type(self, type: Optional[JavaType]) -> TypeHint:
         return self if type is self._type else replace(self, _type=type)
 
     def accept_python(self, v: PythonVisitor[P], p: P) -> J:
@@ -703,9 +703,10 @@ def padding(self) -> PaddingHelper:
         return p
 
     def printer(self, cursor: Cursor) -> TreeVisitor[Tree, PrintOutputCapture[P]]:
-        factory = PrinterFactory.current()
+        if factory := PrinterFactory.current():
+            return factory.create_printer(cursor)
         from .printer import PythonPrinter
-        return factory.create_printer(cursor) if factory else PythonPrinter[PrintOutputCapture[P]]()
+        return PythonPrinter[PrintOutputCapture[P]]()
 
     def accept_python(self, v: PythonVisitor[P], p: P) -> J:
         return v.visit_compilation_unit(self, p)
@@ -1804,13 +1805,13 @@ def expression(self) -> Expression:
     def with_expression(self, expression: Expression) -> YieldFrom:
         return self if expression is self._expression else replace(self, _expression=expression)
 
-    _type: JavaType
+    _type: Optional[JavaType]
 
     @property
-    def type(self) -> JavaType:
+    def type(self) -> Optional[JavaType]:
         return self._type
 
-    def with_type(self, type: JavaType) -> YieldFrom:
+    def with_type(self, type: Optional[JavaType]) -> YieldFrom:
         return self if type is self._type else replace(self, _type=type)
 
     def accept_python(self, v: PythonVisitor[P], p: P) -> J:
@@ -2368,13 +2369,13 @@ def from_(self) -> Expression:
     def with_from(self, from_: Expression) -> ErrorFrom:
         return self.padding.with_from(JLeftPadded.with_element(self._from, from_))
 
-    _type: JavaType
+    _type: Optional[JavaType]
 
     @property
-    def type(self) -> JavaType:
+    def type(self) -> Optional[JavaType]:
         return self._type
 
-    def with_type(self, type: JavaType) -> ErrorFrom:
+    def with_type(self, type: Optional[JavaType]) -> ErrorFrom:
         return self if type is self._type else replace(self, _type=type)
 
     @dataclass