Python: Integrate auto-formatting into template engine (#6728)

Add auto_format() and maybe_auto_format() helper functions to the Python
format module, mirroring the JavaScript template system's pattern. After
template application, _apply_coordinates now calls maybe_auto_format to
normalize spacing, indentation, and other formatting on the templated
subtree using the enclosing CompilationUnit's style context.

The formatter resolves styles from the CompilationUnit in the cursor
ancestry and formats only the templated subtree, not the whole file.
When no CompilationUnit context is available (e.g. synthetic cursors in
unit tests), formatting is gracefully skipped.

Includes rewrite_run integration tests verifying operator spacing,
method call formatting, and correct behavior in deeply nested scopes.

Author: knutwannheden

rewrite-python/rewrite/src/rewrite/python/format/__init__.py

@@ -1,19 +1,40 @@
 __path__ = __import__('pkgutil').extend_path(__path__, __name__)
 
-from .auto_format import AutoFormat
+from typing import Optional
+
+from .auto_format import AutoFormat, AutoFormatVisitor
 from .blank_lines import BlankLinesVisitor
 from .minimum_viable_spacing import MinimumViableSpacingVisitor
 from .normalize_format import NormalizeFormatVisitor
 from .normalize_tabs_or_spaces import NormalizeTabsOrSpacesVisitor
 from .spaces_visitor import SpacesVisitor
 from .tabs_and_indents_visitor import TabsAndIndentsVisitor
 
+from ...visitor import Cursor
+
+
+def auto_format(tree, p, stop_after=None, cursor: Optional[Cursor] = None):
+    try:
+        return AutoFormatVisitor(stop_after).visit(tree, p, cursor)
+    except (ValueError, AttributeError):
+        return tree
+
+
+def maybe_auto_format(before, after, p, stop_after=None, cursor: Optional[Cursor] = None):
+    if before is not after:
+        return auto_format(after, p, stop_after, cursor)
+    return after
+
+
 __all__ = [
     'AutoFormat',
+    'AutoFormatVisitor',
     'BlankLinesVisitor',
     'MinimumViableSpacingVisitor',
     'NormalizeFormatVisitor',
     'NormalizeTabsOrSpacesVisitor',
     'SpacesVisitor',
     'TabsAndIndentsVisitor',
+    'auto_format',
+    'maybe_auto_format',
 ]

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

@@ -251,7 +251,7 @@ def _compare_arguments(
 
         # Check for variadic capture
         if len(pattern_elements) == 1:
-            pattern_arg = pattern_elements[0].element
+            pattern_arg = pattern_elements[0]
             if isinstance(pattern_arg, j.Identifier):
                 cap_name = from_placeholder(pattern_arg.simple_name)
                 if cap_name and self._captures.get(cap_name, Capture(name=cap_name)).variadic:
@@ -267,7 +267,7 @@ def _compare_arguments(
 
         # Compare each argument
         for p_elem, t_elem in zip(pattern_elements, target_elements):
-            if not self._compare(p_elem.element, t_elem.element, cursor):
+            if not self._compare(p_elem, t_elem, cursor):
                 return False
 
         return True

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

@@ -325,6 +325,18 @@ def _apply_coordinates(
                 result,
             )
 
+        # Auto-format the result
+        try:
+            from ..format import maybe_auto_format
+            original = coordinates.tree
+            result = maybe_auto_format(
+                original, result, None, None,
+                cursor.parent if cursor.parent is not None else None
+            )
+        except (ValueError, AttributeError):
+            # No CompilationUnit in cursor ancestry — skip formatting
+            pass
+
         return result
 
     @classmethod

rewrite-python/rewrite/tests/python/format/__init__.py

@@ -0,0 +1,47 @@
+# Copyright 2025 the original author or authors.
+# <p>
+# Licensed under the Moderne Source Available License (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# <p>
+# https://docs.moderne.io/licensing/moderne-source-available-license
+# <p>
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for auto_format and maybe_auto_format helper functions."""
+
+from rewrite.python.format import auto_format, maybe_auto_format
+from rewrite.python.template.engine import TemplateEngine
+
+
+class TestMaybeAutoFormat:
+    """Tests for maybe_auto_format identity guard."""
+
+    def test_same_object_returns_unchanged(self):
+        """When before is after (same identity), return after without formatting."""
+        tree = TemplateEngine.get_template_tree("x+1", {})
+        result = maybe_auto_format(tree, tree, None)
+        assert result is tree
+
+    def test_different_object_no_cursor_skips_gracefully(self):
+        """When before is not after but no cursor, skip formatting without crash."""
+        before = TemplateEngine.get_template_tree("x + 1", {})
+        TemplateEngine.clear_cache()
+        after = TemplateEngine.get_template_tree("y + 2", {})
+
+        result = maybe_auto_format(before, after, None)
+        assert result is not None
+
+
+class TestAutoFormat:
+    """Tests for auto_format fallback behavior."""
+
+    def test_no_cursor_returns_tree_unchanged(self):
+        """auto_format without cursor skips formatting and returns tree."""
+        tree = TemplateEngine.get_template_tree("x+1", {})
+        result = auto_format(tree, None)
+        assert result is tree