Python: Parenthesize template results based on surrounding context

Template.apply() now checks whether the returned expression needs
parentheses when placed in the surrounding AST context, mirroring
JavaTemplate.doApply()'s call to ParenthesizeVisitor.maybeParenthesize.

Restructures Template.apply() into three phases:
1. Placeholder substitution
2. Precedence-based parenthesization (new)
3. Coordinate application (prefix, statement wrapping, formatting)

Author: knutwannheden

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

@@ -139,7 +139,7 @@ def apply_substitutions(
 
         # Apply coordinates (wrap in statement if needed, etc.)
         if coordinates and result is not None:
-            result = cls._apply_coordinates(result, cursor, coordinates)
+            result = cls.apply_coordinates(result, cursor, coordinates)
 
         return result
 
@@ -286,7 +286,7 @@ def _extract_from_method_body(cls, method: j.MethodDeclaration, is_expression: b
         return statements[0]
 
     @classmethod
-    def _apply_coordinates(
+    def apply_coordinates(
         cls,
         result: J,
         cursor: 'Cursor',

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

@@ -27,7 +27,7 @@
 from .placeholder import from_placeholder
 
 if TYPE_CHECKING:
-    pass
+    from rewrite.visitor import Cursor
 
 
 # Operator precedence for Python binary operators (higher number = higher precedence).
@@ -107,6 +107,36 @@ def _needs_parens_under_not(child: J) -> bool:
     return child_prec < 3
 
 
+def maybe_parenthesize(result: J, cursor: 'Cursor') -> J:
+    """Wrap *result* in parentheses when its precedence is lower than the
+    surrounding context (the cursor's parent tree node).
+
+    This mirrors ``ParenthesizeVisitor.maybeParenthesize`` in the Java
+    ``JavaTemplate`` implementation.
+    """
+    if not isinstance(result, (j.Binary, py.Binary, j.Unary)):
+        return result
+
+    parent_cursor = cursor.parent_tree_cursor()
+    parent = parent_cursor.value
+    if not isinstance(parent, J):
+        return result
+
+    # Parent is a binary operator — check precedence
+    if isinstance(parent, (j.Binary, py.Binary)):
+        parent_prec = _get_precedence(parent)
+        result_prec = _get_precedence(result)
+        if parent_prec is not None and result_prec is not None and result_prec < parent_prec:
+            return _wrap_in_parens(result)
+
+    # Parent is `not` — or/and need parens underneath
+    if isinstance(parent, j.Unary) and parent.operator == j.Unary.Type.Not:
+        if _needs_parens_under_not(result):
+            return _wrap_in_parens(result)
+
+    return result
+
+
 class PlaceholderReplacementVisitor(PythonVisitor[None]):
     """
     Visitor that replaces placeholder identifiers with actual values.

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

@@ -138,27 +138,33 @@ def apply(
                 # Assume it's a MatchResult
                 values_dict = values.as_dict()
 
-        # Apply substitutions
+        # Phase 1: placeholder substitution (no coordinates yet)
         if values_dict:
             result = TemplateEngine.apply_substitutions(
                 template_tree,
                 values_dict,
                 cursor,
-                coordinates,
             )
         else:
             result = template_tree
 
-        # If no explicit coordinates, create default replacement coordinates
-        if coordinates is None and cursor is not None:
+        # Phase 2: parenthesize the result if it has lower precedence than
+        # the surrounding context, mirroring JavaTemplate.doApply().
+        # This must happen before coordinates are applied, because
+        # _apply_coordinates may wrap the expression in ExpressionStatement.
+        if result is not None and cursor is not None:
+            from .replacement import maybe_parenthesize
+            result = maybe_parenthesize(result, cursor)
+
+        # Phase 3: apply coordinates (prefix preservation, statement wrapping, auto-format)
+        effective_coords = coordinates
+        if effective_coords is None and cursor is not None:
             tree = cursor.value
             if tree is not None:
-                result = TemplateEngine.apply_substitutions(
-                    result if values_dict else template_tree,
-                    values_dict,
-                    cursor,
-                    PythonCoordinates.replace(tree),
-                )
+                effective_coords = PythonCoordinates.replace(tree)
+
+        if effective_coords is not None and result is not None:
+            result = TemplateEngine.apply_coordinates(result, cursor, effective_coords)
 
         return result
 

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

@@ -14,15 +14,22 @@
 
 """Tests for Template class."""
 
+from typing import Any, Optional
 from uuid import uuid4
 
 import pytest
 
+from rewrite import ExecutionContext, Recipe, TreeVisitor
 from rewrite.java import tree as j
-from rewrite.java.support_types import Space
+from rewrite.java.tree import Unary
+from rewrite.java.support_types import Space, JRightPadded
 from rewrite.markers import Markers
-from rewrite.python.template import template, capture, Template, TemplateBuilder
+from rewrite.python.template import template, capture, pattern, Template, TemplateBuilder
 from rewrite.python.template.engine import TemplateEngine
+from rewrite.python.template.replacement import maybe_parenthesize
+from rewrite.python.visitor import PythonVisitor
+from rewrite.test import RecipeSpec, python
+from rewrite.visitor import Cursor
 
 
 class TestTemplate:
@@ -362,3 +369,258 @@ def test_apply_with_match_result(self):
         assert len(result.arguments) == 1
         assert isinstance(result.arguments[0], j.Identifier)
         assert result.arguments[0].simple_name == "world"
+
+
+# ---------------------------------------------------------------------------
+# maybe_parenthesize unit tests
+# ---------------------------------------------------------------------------
+
+
+class TestMaybeParenthesize:
+    """Unit tests for maybe_parenthesize() — verifies that the function
+    wraps a result node in parentheses when required by the parent context.
+    """
+
+    @staticmethod
+    def _ident(name: str) -> j.Identifier:
+        return j.Identifier(uuid4(), Space.EMPTY, Markers.EMPTY, [], name, None, None)
+
+    @staticmethod
+    def _binary(left, op, right) -> j.Binary:
+        return j.Binary(
+            uuid4(), Space.EMPTY, Markers.EMPTY, left,
+            JRightPadded(op, Space([], ' '), Markers.EMPTY),
+            right, None,
+        )
+
+    @staticmethod
+    def _cursor_chain(*nodes) -> Cursor:
+        """Build a cursor chain from root → … → leaf."""
+        cur = Cursor(None, Cursor.ROOT_VALUE)
+        for node in nodes:
+            cur = Cursor(cur, node)
+        return cur
+
+    def test_or_result_in_and_parent(self):
+        """An `or` expression placed inside an `and` must be parenthesized."""
+        or_expr = self._binary(self._ident('a'), j.Binary.Type.Or, self._ident('b'))
+        and_parent = self._binary(or_expr, j.Binary.Type.And, self._ident('z'))
+        cursor = self._cursor_chain(and_parent, or_expr)
+
+        result = maybe_parenthesize(or_expr, cursor)
+        assert isinstance(result, j.Parentheses)
+
+    def test_and_result_in_or_parent(self):
+        """`and` has higher precedence than `or`, so no parens needed."""
+        and_expr = self._binary(self._ident('a'), j.Binary.Type.And, self._ident('b'))
+        or_parent = self._binary(and_expr, j.Binary.Type.Or, self._ident('z'))
+        cursor = self._cursor_chain(or_parent, and_expr)
+
+        result = maybe_parenthesize(and_expr, cursor)
+        assert not isinstance(result, j.Parentheses)
+
+    def test_or_result_under_not(self):
+        """`or` under `not` needs parens (not has precedence 3, or is 1)."""
+        or_expr = self._binary(self._ident('a'), j.Binary.Type.Or, self._ident('b'))
+        not_parent = j.Unary(
+            uuid4(), Space.EMPTY, Markers.EMPTY,
+            JRightPadded(j.Unary.Type.Not, Space.EMPTY, Markers.EMPTY),
+            or_expr, None,
+        )
+        cursor = self._cursor_chain(not_parent, or_expr)
+
+        result = maybe_parenthesize(or_expr, cursor)
+        assert isinstance(result, j.Parentheses)
+
+    def test_identifier_result_unchanged(self):
+        """Non-binary/unary results are returned unchanged."""
+        ident = self._ident('x')
+        and_parent = self._binary(ident, j.Binary.Type.And, self._ident('z'))
+        cursor = self._cursor_chain(and_parent, ident)
+
+        result = maybe_parenthesize(ident, cursor)
+        assert result is ident
+
+    def test_same_precedence_no_parens(self):
+        """`and` inside `and` — same precedence, no parens needed."""
+        inner = self._binary(self._ident('a'), j.Binary.Type.And, self._ident('b'))
+        outer = self._binary(inner, j.Binary.Type.And, self._ident('c'))
+        cursor = self._cursor_chain(outer, inner)
+
+        result = maybe_parenthesize(inner, cursor)
+        assert not isinstance(result, j.Parentheses)
+
+    def test_template_apply_parenthesizes_in_binary_context(self):
+        """template.apply() should parenthesize its result when the cursor
+        parent is a higher-precedence binary operator."""
+        from rewrite.python.tree import ExpressionStatement
+
+        _x = capture('x')
+        _y = capture('y')
+        _not_x_or_not_y = template("not {x} or not {y}", x=_x, y=_y)
+
+        # Build cursor: root → Binary(and) → Unary(not)
+        # The Unary is the node being replaced by the template result.
+        x_ident = self._ident('x')
+        y_ident = self._ident('y')
+        and_inner = self._binary(x_ident, j.Binary.Type.And, y_ident)
+        not_expr = j.Unary(
+            uuid4(), Space.EMPTY, Markers.EMPTY,
+            JRightPadded(j.Unary.Type.Not, Space.EMPTY, Markers.EMPTY),
+            j.Parentheses(
+                uuid4(), Space.EMPTY, Markers.EMPTY,
+                JRightPadded(and_inner, Space.EMPTY, Markers.EMPTY),
+            ),
+            None,
+        )
+        z_ident = self._ident('z')
+        outer_and = self._binary(not_expr, j.Binary.Type.And, z_ident)
+        cursor = self._cursor_chain(outer_and, not_expr)
+
+        result = _not_x_or_not_y.apply(
+            cursor, values={'x': x_ident, 'y': y_ident}
+        )
+
+        # apply() wraps in ExpressionStatement (target Unary is a Statement);
+        # the inner expression should be Parentheses since `or` is placed
+        # inside an `and` context.
+        inner = result.expression if isinstance(result, ExpressionStatement) else result
+        assert isinstance(inner, j.Parentheses), (
+            f"Expected Parentheses but got {type(inner).__name__}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Template apply precedence in surrounding context (integration tests)
+# ---------------------------------------------------------------------------
+#
+# When template.apply() returns a node with lower precedence than the
+# site where it replaces, the framework should wrap it in parentheses.
+#
+# Example: "not {x} or not {y}" produces an `or` expression. If the
+# replacement site's parent is `and`, the result needs outer parentheses
+# because `or` has lower precedence than `and`.
+# ---------------------------------------------------------------------------
+
+
+class _DeMorganRecipe(Recipe):
+    """Minimal De Morgan recipe that demonstrates the parenthesization bug."""
+
+    @property
+    def name(self) -> str:
+        return "test.DeMorgan"
+
+    @property
+    def display_name(self) -> str:
+        return "Test De Morgan"
+
+    @property
+    def description(self) -> str:
+        return "Test De Morgan parenthesization"
+
+    def editor(self) -> TreeVisitor[Any, ExecutionContext]:
+        _x = capture('x')
+        _y = capture('y')
+        _not_and = pattern("not ({x} and {y})", x=_x, y=_y)
+        _not_or = pattern("not ({x} or {y})", x=_x, y=_y)
+        _not_x_or_not_y = template("not {x} or not {y}", x=_x, y=_y)
+        _not_x_and_not_y = template("not {x} and not {y}", x=_x, y=_y)
+
+        class Visitor(PythonVisitor[ExecutionContext]):
+            def visit_unary(
+                self, unary: Unary, p: ExecutionContext
+            ) -> Optional[Unary]:
+                unary = super().visit_unary(unary, p)
+
+                match = _not_and.match(unary, self.cursor)
+                if match:
+                    return _not_x_or_not_y.apply(self.cursor, values=match)
+
+                match = _not_or.match(unary, self.cursor)
+                if match:
+                    return _not_x_and_not_y.apply(self.cursor, values=match)
+
+                return unary
+
+        return Visitor()
+
+
+class TestTemplateApplyPrecedenceInContext:
+    """Template.apply() result must be parenthesized when inserted into
+    a higher-precedence context.
+
+    All tests use a minimal De Morgan recipe:
+        not (x and y) → not x or not y
+        not (x or y)  → not x and not y
+    """
+
+    def test_demorgan_standalone(self):
+        """Baseline: De Morgan works at top level (no surrounding context)."""
+        spec = RecipeSpec(recipe=_DeMorganRecipe())
+        spec.rewrite_run(
+            python(
+                "result = not (a and b)",
+                "result = not a or not b",
+            )
+        )
+
+    def test_demorgan_result_in_and_context(self):
+        """not (x and y) inside `... and z` must wrap result in parens.
+
+        not (x and y) and z
+          → (not x or not y) and z
+
+        Without parens: not x or not y and z
+        Python parses:  (not x) or ((not y) and z)  ← WRONG
+        """
+        spec = RecipeSpec(recipe=_DeMorganRecipe())
+        spec.rewrite_run(
+            python(
+                "result = not (x and y) and z",
+                "result = (not x or not y) and z",
+            )
+        )
+
+    def test_demorgan_result_in_or_context(self):
+        """not (x or y) inside `... or z` — no parens needed.
+
+        not (x or y) or z
+          → not x and not y or z
+
+        `and` has higher precedence than `or`, so
+        Python parses:  ((not x) and (not y)) or z  ← correct without parens.
+        """
+        spec = RecipeSpec(recipe=_DeMorganRecipe())
+        spec.rewrite_run(
+            python(
+                "result = not (x or y) or z",
+                "result = not x and not y or z",
+            )
+        )
+
+    def test_demorgan_result_on_right_of_and(self):
+        """z and not (x and y) must wrap result in parens.
+
+        z and not (x and y)
+          → z and (not x or not y)
+
+        Without parens: z and not x or not y
+        Python parses:  (z and (not x)) or (not y)  ← WRONG
+        """
+        spec = RecipeSpec(recipe=_DeMorganRecipe())
+        spec.rewrite_run(
+            python(
+                "result = z and not (x and y)",
+                "result = z and (not x or not y)",
+            )
+        )
+
+    def test_demorgan_result_in_if_condition(self):
+        """De Morgan in an if condition should still be correct."""
+        spec = RecipeSpec(recipe=_DeMorganRecipe())
+        spec.rewrite_run(
+            python(
+                "if not (x and y) and z:\n    pass",
+                "if (not x or not y) and z:\n    pass",
+            )
+        )