fix super replacement of multiple lines

asottile · Sep 25, 2021 · dd1fe7f · dd1fe7f
1 parent eadaf86
commit dd1fe7f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 20 deletions.
diff --git a/pyupgrade/_plugins/legacy.py b/pyupgrade/_plugins/legacy.py
@@ -21,9 +21,7 @@
 from pyupgrade._token_helpers import Block
 from pyupgrade._token_helpers import find_and_replace_call
 from pyupgrade._token_helpers import find_block_start
-from pyupgrade._token_helpers import find_open_paren
 from pyupgrade._token_helpers import find_token
-from pyupgrade._token_helpers import parse_call_args
 
 FUNC_TYPES = (ast.Lambda, ast.FunctionDef, ast.AsyncFunctionDef)
 NON_LAMBDA_FUNC_TYPES = (ast.FunctionDef, ast.AsyncFunctionDef)
@@ -37,20 +35,6 @@ def _fix_yield(i: int, tokens: List[Token]) -> None:
     tokens[i:block.end] = [Token('CODE', f'yield from {container}\n')]
 
 
-def _fix_old_super(i: int, tokens: List[Token]) -> None:
-    j = find_open_paren(tokens, i)
-    k = j - 1
-    while tokens[k].src != '.':
-        k -= 1
-    func_args, end = parse_call_args(tokens, j)
-    # remove the first argument
-    if len(func_args) == 1:
-        del tokens[func_args[0][0]:func_args[0][0] + 1]
-    else:
-        del tokens[func_args[0][0]:func_args[1][0] + 1]
-    tokens[i:k] = [Token('CODE', 'super()')]
-
-
 def _is_simple_base(base: ast.AST) -> bool:
     return (
         isinstance(base, ast.Name) or (
@@ -76,7 +60,7 @@ class Visitor(ast.NodeVisitor):
     def __init__(self) -> None:
         self._scopes: List[Scope] = []
         self.super_offsets: Set[Offset] = set()
-        self.old_super_offsets: Set[Offset] = set()
+        self.old_super_offsets: Set[Tuple[Offset, str]] = set()
         self.yield_offsets: Set[Offset] = set()
 
     @contextlib.contextmanager
@@ -153,7 +137,7 @@ def visit_Call(self, node: ast.Call) -> None:
                     node.func.value,
                 )
         ):
-            self.old_super_offsets.add(ast_to_offset(node))
+            self.old_super_offsets.add((ast_to_offset(node), node.func.attr))
 
         self.generic_visit(node)
 
@@ -204,8 +188,10 @@ def visit_Module(
     for offset in visitor.super_offsets:
         yield offset, super_func
 
-    for offset in visitor.old_super_offsets:
-        yield offset, _fix_old_super
+    for offset, func_name in visitor.old_super_offsets:
+        template = f'super().{func_name}({{rest}})'
+        callback = functools.partial(find_and_replace_call, template=template)
+        yield offset, callback
 
     for offset in visitor.yield_offsets:
         yield offset, _fix_yield
diff --git a/tests/features/super_test.py b/tests/features/super_test.py
@@ -201,6 +201,21 @@ def test_old_style_class_super_noop(s):
             '        super().f()\n'
             '        super().f(arg, arg)\n',
         ),
+        pytest.param(
+            'class C(B):\n'
+            '    def f(self, a):\n'
+            '         B.f(\n'
+            '             self,\n'
+            '             a,\n'
+            '         )\n',
+
+            'class C(B):\n'
+            '    def f(self, a):\n'
+            '         super().f(\n'
+            '             a,\n'
+            '         )\n',
+            id='multi-line super call',
+        ),
     ),
 )
 def test_old_style_class_super(s, expected):