Python - fix for method_select in TabsAndIndentsVisitor (#6711)

* Fix for method_select

* Clearer impl of _compute_select_column

Author: greg-at-moderne

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

@@ -215,8 +215,12 @@ def visit_right_padded(self, right: Optional[JRightPadded], p: P) -> Optional[JR
                     after = self._indent_to(right.after, indent)
 
                 else:
+                    method_select_indent = self.cursor.get_nearest_message("method_select_indent")
                     elem = self.visit_and_cast(elem, J, p)
-                    after = self.visit_space(right.after, p)
+                    if method_select_indent is not None:
+                        after = self._indent_to(right.after, method_select_indent)
+                    else:
+                        after = self.visit_space(right.after, p)
 
             else:
                 if rp_context in ("method_invocation_argument",):
@@ -425,6 +429,62 @@ def visit_class_declaration(self, class_decl: ClassDeclaration, p: P) -> J:
             class_decl.body, Block, p))
         return class_decl
 
+    def visit_method_invocation(self, method: MethodInvocation, p: P) -> J:
+        select = method.padding.select
+        if select is not None and '\n' in select.after.last_whitespace:
+            col = self._compute_select_column(method)
+            if col >= 0:
+                self.cursor.put_message("method_select_indent", col)
+        return super().visit_method_invocation(method, p)
+
+    def _compute_select_column(self, method: MethodInvocation) -> int:
+        from rewrite.python.printer import PythonPrinter, PrintOutputCapture
+
+        line_start = None
+        for c in self.cursor.get_path_as_cursors():
+            v = c.value
+            if isinstance(v, J):
+                line_start = v
+                if '\n' in v.prefix.whitespace:
+                    break
+        if line_start is None:
+            return -1
+
+        select = method.select
+
+        class _ColumnCounter(PrintOutputCapture):
+            def __init__(self):
+                super().__init__()
+                self.col = 0
+                self.found = False
+
+            def append(self, text):
+                if text and not self.found:
+                    for ch in text:
+                        self.col = 0 if ch == '\n' else self.col + 1
+                return self
+
+        class _Printer(PythonPrinter):
+            def __init__(self, target):
+                super().__init__()
+                orig_visit = self._delegate.visit
+                def _check(tree, p):
+                    if tree is target:
+                        p.found = True
+                        return tree
+                    return orig_visit(tree, p) if not p.found else tree
+                self._delegate.visit = _check
+
+            def visit(self, tree, p, parent=None):
+                if p.found or tree is select:
+                    p.found = True
+                    return tree
+                return super().visit(tree, p)
+
+        counter = _ColumnCounter()
+        _Printer(select).print(line_start, counter)
+        return counter.col if counter.found else -1
+
     # -------------------------------------------------------------------------
     # Expression statement (docstring alignment)
     # -------------------------------------------------------------------------

rewrite-python/rewrite/tests/python/all/format/tabs_and_indents_visitor_test.py

@@ -994,7 +994,6 @@ def test_top_level_string():
     )
 
 
-@pytest.mark.xfail
 def test_method_select_suffix_new_line_already_correct():
     style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4)
     rewrite_run(
@@ -1008,7 +1007,6 @@ def test_method_select_suffix_new_line_already_correct():
     )
 
 
-@pytest.mark.xfail
 def test_method_select_suffix():
     style = IntelliJ.tabs_and_indents().with_use_tab_character(False).with_tab_size(4)
     rewrite_run(