diff --git a/rewrite-python/rewrite/src/rewrite/python/template/replacement.py b/rewrite-python/rewrite/src/rewrite/python/template/replacement.py index 5a71ec24d38..124829f2984 100644 --- a/rewrite-python/rewrite/src/rewrite/python/template/replacement.py +++ b/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) diff --git a/rewrite-python/rewrite/tests/python/template/test_replacement.py b/rewrite-python/rewrite/tests/python/template/test_replacement.py index 9112c1e4af8..a3bae66f902 100644 --- a/rewrite-python/rewrite/tests/python/template/test_replacement.py +++ b/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__}" + )