Python: Unwrap ExpressionStatement around block-level template placeholders (#6838)

When a template has a placeholder in statement position (e.g. the body
of a `with` block), the parser wraps it as ExpressionStatement(Identifier).
If the replacement is a non-Expression statement (return, if, for, etc.),
the visitor must substitute it directly in the block before the default
visitor runs, so visit_expression_statement never tries to cast the
replacement to Expression. This mirrors the JS template engine's
visitBlock logic.

Author: knutwannheden

rewrite-python/rewrite/src/rewrite/python/template/replacement.py

@@ -193,6 +193,42 @@ def visit_identifier(self, ident: j.Identifier, p: None) -> J:
         # Not a placeholder or no value provided, continue normally
         return super().visit_identifier(ident, p)
 
+    def visit_block(self, block: j.Block, p: None) -> J:
+        """Override visit_block to unwrap ExpressionStatement around placeholders.
+
+        When a template has a placeholder in statement position (e.g., the body
+        of a ``with`` block), the parser wraps it as ``ExpressionStatement(Identifier(...))``
+        since it looks like a bare expression.  If the replacement value is a
+        non-Expression statement (``return``, ``if``, ``for``, etc.), we must
+        substitute it directly into the statements list, bypassing the
+        ``ExpressionStatement`` wrapper.  This mirrors the JS template engine's
+        ``visitBlock`` logic.
+        """
+        # Substitute statement-position placeholders BEFORE the default
+        # visitor runs, so visit_expression_statement never sees them.
+        new_stmts: List[JRightPadded] = []
+        changed = False
+        for rp in block.padding.statements:
+            stmt = rp.element
+            if isinstance(stmt, py.ExpressionStatement):
+                expr = stmt.expression
+                if isinstance(expr, j.Identifier):
+                    capture_name = from_placeholder(expr.simple_name)
+                    if capture_name is not None and capture_name in self._values:
+                        replacement = self._values[capture_name]
+                        # Preserve the placeholder's whitespace prefix
+                        if hasattr(replacement, '_prefix'):
+                            replacement = replacement.replace(_prefix=expr.prefix)
+                        new_stmts.append(rp.replace(element=replacement))
+                        changed = True
+                        continue
+            new_stmts.append(rp)
+
+        if changed:
+            block = block.padding.replace(statements=new_stmts)
+
+        return super().visit_block(block, p)
+
     def visit_binary(self, binary: j.Binary, p: None) -> J:
         """Visit a Java Binary and auto-parenthesize substituted operands if needed."""
         binary = super().visit_binary(binary, p)

rewrite-python/rewrite/tests/python/template/test_replacement.py

@@ -341,3 +341,39 @@ def test_comparison_under_not_no_parens(self):
         assert isinstance(result, j.Unary)
         # Comparisons have precedence 4 which is >= 3 (not threshold), so no parens
         assert isinstance(result.expression, j.Binary)
+
+
+class TestBlockStatementPlaceholder:
+    """Tests for placeholder replacement in statement position inside blocks.
+
+    When a template has {body} inside a block (e.g. ``if True:\\n    {body}``),
+    the parser wraps it as ``ExpressionStatement(Identifier(placeholder))``.
+    If the replacement is a non-Expression statement (Return, If, etc.),
+    the visitor must unwrap the ExpressionStatement so the block directly
+    contains the replacement statement.
+    """
+
+    def setup_method(self):
+        TemplateEngine.clear_cache()
+
+    def test_return_in_block_placeholder(self):
+        """Substituting a Return for a block placeholder should unwrap ExpressionStatement."""
+        tree = TemplateEngine.get_template_tree(
+            "if True:\n    {body}", {'body': capture('body')}
+        )
+        return_stmt = j.Return(
+            uuid4(), Space([], '    '), Markers.EMPTY, _ident('result')
+        )
+        visitor = PlaceholderReplacementVisitor({'body': return_stmt})
+        result = visitor.visit(tree, None)
+
+        # The result should be an If statement
+        assert isinstance(result, j.If)
+        # The then-part block should contain a Return, NOT an ExpressionStatement
+        then_block = result.then_part
+        assert isinstance(then_block, j.Block)
+        stmts = then_block.statements
+        assert len(stmts) == 1
+        assert isinstance(stmts[0], j.Return), (
+            f"Expected Return in block, got {type(stmts[0]).__name__}"
+        )