Support first_statement coordinate for J.Block

Author: knutwannheden

rewrite/rewrite/java/support_types.py

@@ -359,6 +359,11 @@ def replace(self, loc: Optional[Space.Location] = None) -> JavaCoordinates:
 
 @dataclass
 class _BlockCoordinateBuilder(_StatementCoordinateBuilder):
+    def first_statement(self) -> JavaCoordinates:
+        if not self.tree.statements:
+            return self.last_statement()
+        return self.tree.statements[0].get_coordinates().before()
+
     def last_statement(self) -> JavaCoordinates:
         return self.before(Space.Location.BLOCK_END)
 

rewrite/rewrite/utils.py

@@ -55,9 +55,10 @@ def list_flat_map(fn: FlatMapFnType[T], lst: List[T]) -> List[T]:
         if isinstance(new_items, list) and (len(new_items) != 1 or new_items[0] is not item):
             changed = True
             result.extend(new_items)
-        elif not isinstance(new_items, list) and new_items is not item:
+        elif not isinstance(new_items, list):
+            if changed or new_items is not item:
+                result.append(new_items)
             changed = True
-            result.append(new_items)
 
     return result if changed else lst
 

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

@@ -82,6 +82,32 @@ def f():
     )
 
 
+def test_add_statement_first():
+    rewrite_run(
+        # language=python
+        python(
+            """\
+            def f():
+                pass
+                pass
+            """,
+            """\
+            def f():
+                return
+                pass
+                pass
+            """
+        ),
+        spec=RecipeSpec()
+        .with_recipe(from_visitor(
+            GenericTemplatingVisitor(
+                lambda j: isinstance(j, MethodDeclaration) and len(j.body.statements) == 2,
+                'return',
+                coordinate_provider=lambda m: cast(MethodDeclaration, m).body.get_coordinates().first_statement())
+        ))
+    )
+
+
 def test_add_statement_before():
     rewrite_run(
         # language=python