Support appending statements to blocks using templating

Author: knutwannheden

rewrite/rewrite/java/support_types.py

@@ -335,12 +335,24 @@ def replace(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
 
 @dataclass
 class _ExpressionCoordinateBuilder(CoordinateBuilder):
+    def after(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
+        return JavaCoordinates(self.tree, loc or Space.Location.EXPRESSION_PREFIX, JavaCoordinates.Mode.AFTER)
+
+    def before(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
+        return JavaCoordinates(self.tree, loc or Space.Location.EXPRESSION_PREFIX, JavaCoordinates.Mode.BEFORE)
+
     def replace(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
         return JavaCoordinates(self.tree, loc or Space.Location.EXPRESSION_PREFIX, JavaCoordinates.Mode.REPLACE)
 
 
 @dataclass
 class _StatementCoordinateBuilder(CoordinateBuilder):
+    def after(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
+        return JavaCoordinates(self.tree, loc or Space.Location.STATEMENT_PREFIX, JavaCoordinates.Mode.AFTER)
+
+    def before(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
+        return JavaCoordinates(self.tree, loc or Space.Location.STATEMENT_PREFIX, JavaCoordinates.Mode.BEFORE)
+
     def replace(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
         return JavaCoordinates(self.tree, loc or Space.Location.STATEMENT_PREFIX, JavaCoordinates.Mode.REPLACE)
 

rewrite/rewrite/python/format/auto_format.py

@@ -21,7 +21,7 @@ def get_visitor(self):
         return AutoFormatVisitor()
 
 
-class AutoFormatVisitor(PythonVisitor):
+class AutoFormatVisitor(PythonVisitor[P]):
     def __init__(self, stop_after: Optional[Tree] = None):
         self._stop_after = stop_after
 

rewrite/rewrite/python/templating.py

@@ -2,14 +2,15 @@
 
 import re
 from dataclasses import dataclass, field
-from typing import Any, Callable, Optional, cast, List
+from typing import Any, Callable, Optional, cast, List, Type
 
 from . import CompilationUnit, ExpressionStatement, PythonVisitor
 from .parser import PythonParserBuilder
 from .printer import PythonPrinter
-from .. import PrintOutputCapture
-from ..java import J, JavaCoordinates, Space, Identifier, P, Literal
+from .. import PrintOutputCapture, Tree
+from ..java import J, JavaCoordinates, Space, Identifier, P, Literal, Expression, Statement, Block
 from ..parser import ParserBuilder
+from ..utils import list_flat_map
 from ..visitor import Cursor
 
 
@@ -30,8 +31,9 @@ def apply(self, scope: Cursor, coordinates: JavaCoordinates, *parameters) -> J:
         substituted = substitutions.substitute()
         if self.on_after_variable_substitution:
             self.on_after_variable_substitution(substituted)
-        parsed = self._template_parser.parse_expression(scope, substituted, coordinates.loc)
-        return substitutions.unsubstitute(parsed).with_prefix(cast(J, scope.value).prefix)
+
+        return PythonTemplatePythonExtension(self._template_parser, substitutions, substituted, coordinates)\
+            .visit(cast(Tree, scope.value), 0, scope.parent_or_throw)
 
     def substitutions(self, parameters: List[Any]) -> Substitutions:
         return Substitutions(self.code, parameters)
@@ -50,6 +52,10 @@ def parse_expression(self, scope: Cursor, template: str, loc: Space.Location) ->
             cu.statements[0]
         return j.with_prefix(cast(J, scope.value).prefix)
 
+    def parse_block_statements(self, cursor: Cursor, expected: Type, template: str, loc: Space.Location, mode: JavaCoordinates.Mode) -> List[J]:
+        cu: CompilationUnit = next(iter(self.parser_builder.build().parse_strings(template)))
+        return cu.statements
+
 
 @dataclass
 class Substitutions:
@@ -110,6 +116,9 @@ def substitute_untyped(self, index: int) -> str:
     def unsubstitute(self, parsed: J) -> J:
         return cast(J, UnsubstitutionVisitor(self.parameters).visit(parsed, 0))
 
+    def unsubstitute_all(self, parsed: List[J]) -> List[J]:
+        return [self.unsubstitute(j) for j in parsed]
+
 
 @dataclass
 class UnsubstitutionVisitor(PythonVisitor[int]):
@@ -120,3 +129,46 @@ def visit_identifier(self, identifier: Identifier, p: P) -> J:
         if match := self._param_pattern.fullmatch(identifier.simple_name):
             return cast(J, self.parameters[int(match.group(1))]).with_prefix(identifier.prefix)
         return identifier
+
+@dataclass
+class PythonTemplatePythonExtension(PythonVisitor[int]):
+    template_parser: PythonTemplateParser
+    substitutions: Substitutions
+    substituted_template: str
+    coordinates: JavaCoordinates
+
+    def __post_init__(self):
+        self.insertion_point = self.coordinates.tree
+        self.loc = self.coordinates.loc
+        self.mode = self.coordinates.mode
+
+    def visit_expression(self, expression: Expression, p: int) -> J:
+        if (self.loc == Space.Location.EXPRESSION_PREFIX or self.loc == Space.Location.STATEMENT_PREFIX and
+            isinstance(expression, Statement)) and expression.is_scope(self.insertion_point):
+            parsed = self.template_parser.parse_expression(self.cursor, self.substituted_template, self.loc)
+            return self.auto_format(self.substitutions.unsubstitute(parsed).with_prefix(expression.prefix), p)
+        return expression
+
+    def visit_block(self, block: Block, p: P) -> J:
+        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)
+
+    def visit_statement(self, statement: Statement, p: P) -> J:
+        return statement
+        # if (self.loc == Space.Location.STATEMENT_PREFIX and statement.is_scope(self.insertion_point):
+        #     parsed = self.template_parser.parse_expression(self.cursor, self.substituted_template, self.loc)
+        #     return self.auto_format(self.substitutions.unsubstitute(parsed).with_prefix(expression.prefix), p)
+        # return expression
+
+    def get_replacements(self, statement: Statement) -> List[J]:
+        parsed = self.template_parser.parse_block_statements(Cursor(self.cursor, self.insertion_point), Statement, self.substituted_template, self.loc, self.mode)
+        gen = self.substitutions.unsubstitute_all(parsed)
+        formatted = [s.with_prefix(statement.prefix.with_comments([])) for s in gen]
+
+        if self.mode == JavaCoordinates.Mode.REPLACE:
+            return formatted
+        elif self.mode == JavaCoordinates.Mode.BEFORE:
+            return formatted + [statement]
+        else:
+            return [statement] + formatted

rewrite/rewrite/python/visitor.py

@@ -11,6 +11,11 @@ class PythonVisitor(JavaVisitor[P]):
     def is_acceptable(self, source_file: SourceFile, p: P) -> bool:
         return isinstance(source_file, Py)
 
+    def auto_format(self, j: J, p: P, cursor: Optional[Cursor] = None, stop_after: Optional[J] = None) -> J:
+        cursor = cursor or self.cursor.parent_tree_cursor()
+        from .format import AutoFormatVisitor
+        return AutoFormatVisitor(stop_after).visit(j, p, cursor)
+
     def visit_async(self, async_: Async, p: P) -> J:
         async_ = async_.with_prefix(self.visit_space(async_.prefix, PySpace.Location.ASYNC_PREFIX, p))
         temp_statement = cast(Statement, self.visit_statement(async_, p))

rewrite/rewrite/tree.py

@@ -54,6 +54,9 @@ def print(self, cursor: 'Cursor', capture: 'PrintOutputCapture[P]') -> str:
     def printer(self, cursor: 'Cursor') -> 'TreeVisitor[Any, PrintOutputCapture[P]]':
         return cursor.first_enclosing_or_throw(SourceFile).printer(cursor)
 
+    def is_scope(self, tree: Optional[Tree]) -> bool:
+        return tree and tree.id == self.id
+
     def __eq__(self, other: object) -> bool:
         if self.__class__ == other.__class__:
             return self.id == cast(Tree, other).id

rewrite/rewrite/utils.py

@@ -10,6 +10,7 @@ def random_id() -> UUID:
 
 # Define a type that allows both single and two-argument callables
 FnType = Union[Callable[[T], Union[T, None]], Callable[[T, int], Union[T, None]]]
+FlatMapFnType = Union[Callable[[T], List[T]], Callable[[T, int], List[T]]]
 
 def list_map(fn: FnType[T], lst: List[T]) -> List[T]:
     changed = False
@@ -33,6 +34,24 @@ def list_map(fn: FnType[T], lst: List[T]) -> List[T]:
     return mapped_lst if changed else lst  # type: ignore
 
 
+def list_flat_map(fn: FlatMapFnType[T], lst: List[T]) -> List[T]:
+    changed = False
+    result: List[T] = []
+
+    with_index = len(inspect.signature(fn).parameters) == 2
+    for index, item in enumerate(lst):
+        new_items = fn(item, index) if with_index else fn(item)  # type: ignore
+        if new_items is None:
+            changed = True
+            continue
+
+        if len(new_items) != 1 or new_items[0] is not item:
+            changed = True
+        result.extend(new_items)
+
+    return result if changed else lst
+
+
 def list_map_last(fn: Callable[[T], Union[T, None]], lst: List[T]) -> List[T]:
     if not lst:
         return lst

rewrite/tests/python/all/templating/template_test.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
-from typing import Any, Callable, List, cast
+from typing import Any, Callable, List, cast, Union
 
-from rewrite.java import Literal, P, J, Expression
+from rewrite.java import Literal, P, J, Expression, Statement, JavaCoordinates, MethodDeclaration
 from rewrite.python import PythonVisitor, PythonTemplate, PythonParserBuilder, CompilationUnit, ExpressionStatement
 from rewrite.test import from_visitor, RecipeSpec, rewrite_run, python
 
@@ -15,7 +15,7 @@ def test_string_substitution():
         ),
         spec=RecipeSpec()
         .with_recipe(from_visitor(
-            ExpressionTemplatingVisitor(lambda j: isinstance(j, Literal), '#{}', [2])))
+            ReplaceTemplatingVisitor(lambda j: isinstance(j, Literal), '#{}', [2])))
     )
 
 
@@ -28,7 +28,31 @@ def test_tree_substitution():
         ),
         spec=RecipeSpec()
         .with_recipe(from_visitor(
-            ExpressionTemplatingVisitor(lambda j: isinstance(j, Literal), '#{any()}', [parse_expression('2')])))
+            ReplaceTemplatingVisitor(lambda j: isinstance(j, Literal), '#{any()}', [parse_expression('2')])))
+    )
+
+
+def test_add_statement():
+    rewrite_run(
+        # language=python
+        python(
+            """\
+            def f():
+                pass
+            """,
+            """\
+            def f():
+                pass
+                return
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipe(from_visitor(
+            AddLastTemplatingVisitor(
+                lambda j: isinstance(j, MethodDeclaration) and len(j.body.statements) == 1,
+                'return',
+                coordinate_provider=lambda m: cast(MethodDeclaration, m).body.statements[0].get_coordinates().after())
+        ))
     )
 
 
@@ -41,18 +65,20 @@ def test_tree_substitution_named():
         ),
         spec=RecipeSpec()
         .with_recipe(from_visitor(
-            ExpressionTemplatingVisitor(lambda j: isinstance(j, Literal), '#{name:any()} + #{any()} + #{name}', [parse_expression('2'), parse_expression('3')])))
+            ReplaceTemplatingVisitor(lambda j: isinstance(j, Literal), '#{name:any()} + #{any()} + #{name}', [parse_expression('2'), parse_expression('3')])))
     )
 
 
 def parse_expression(code: str) -> J:
-    return cast(ExpressionStatement,
-                cast(CompilationUnit, next(iter(PythonParserBuilder().build().parse_strings(code)))).statements[
-                    0]).expression
+    return cast(ExpressionStatement, parse_statement(code)).expression
+
+
+def parse_statement(code: str) -> J:
+    return cast(CompilationUnit, next(iter(PythonParserBuilder().build().parse_strings(code)))).statements[0]
 
 
 @dataclass
-class ExpressionTemplatingVisitor(PythonVisitor[P]):
+class ReplaceTemplatingVisitor(PythonVisitor[P]):
     match: Callable[[J], bool]
     code: str
     params: List[Any] = field(default_factory=list)
@@ -68,3 +94,32 @@ def __post_init__(self):
     def visit_expression(self, expr: Expression, p: P) -> J:
         return self._template.apply(self.cursor, expr.get_coordinates().replace(), self.params) if self.match(
             expr) else expr
+
+    def visit_statement(self, stmt: Statement, p: P) -> J:
+        return self._template.apply(self.cursor, stmt.get_coordinates().replace(), self.params) if self.match(
+            stmt) else stmt
+
+
+@dataclass
+class AddLastTemplatingVisitor(PythonVisitor[P]):
+    match: Callable[[J], bool]
+    code: str
+    coordinate_provider: Callable[[Union[Expression, Statement]], JavaCoordinates] = lambda j: j.get_coordinates().after()
+    params: List[Any] = field(default_factory=list)
+
+    debug: bool = False
+    _template: PythonTemplate = field(init=False, repr=False)
+
+    def __post_init__(self):
+        self._template = PythonTemplate(
+            self.code,
+            on_after_variable_substitution=lambda code: print(code) if self.debug else None
+        )
+
+    def visit_expression(self, expr: Expression, p: P) -> J:
+        return self._template.apply(self.cursor, self.coordinate_provider(expr), self.params) if self.match(
+            expr) else expr
+
+    def visit_statement(self, stmt: Statement, p: P) -> J:
+        return self._template.apply(self.cursor, self.coordinate_provider(stmt), self.params) if self.match(
+            stmt) else stmt