Skip to content

Commit 00aa729

Browse files
hamdanalasottile
hamdanal
authored and
asottile
committed
rewrite 2-arg super call in nested class
1 parent d9461f5 commit 00aa729

2 files changed

Lines changed: 66 additions & 6 deletions

File tree

‎pyupgrade/_plugins/legacy.py‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -131,19 +131,42 @@ def visit_Call(self, node: ast.Call) -> None:
131131
isinstance(node.func, ast.Name) and
132132
node.func.id == 'super' and
133133
len(node.args) == 2 and
134-
isinstance(node.args[0], ast.Name) and
135134
isinstance(node.args[1], ast.Name) and
136135
# there are at least two scopes
137136
len(self._scopes) >= 2 and
138-
# the second to last scope is the class in arg1
139-
isinstance(self._scopes[-2].node, ast.ClassDef) and
140-
node.args[0].id == self._scopes[-2].node.name and
141137
# the last scope is a function where the first arg is arg2
142138
isinstance(self._scopes[-1].node, FUNC_TYPES) and
143139
self._scopes[-1].node.args.args and
144140
node.args[1].id == self._scopes[-1].node.args.args[0].arg
145141
):
146-
self.super_offsets.add(ast_to_offset(node))
142+
args = node.args[0]
143+
scope = len(self._scopes) - 2
144+
current_scope = self._scopes[scope]
145+
# if in nested classes, all names in arg1 must match the scopes
146+
while (
147+
isinstance(args, ast.Attribute) and
148+
scope > 0 and
149+
isinstance(current_scope.node, ast.ClassDef) and
150+
args.attr == current_scope.node.name
151+
):
152+
args = args.value
153+
scope -= 1
154+
current_scope = self._scopes[scope]
155+
# now check if it is outer most class and its name match
156+
if (
157+
isinstance(args, ast.Name) and
158+
isinstance(current_scope.node, ast.ClassDef) and
159+
args.id == current_scope.node.name and
160+
# an enclosing scope cannot be a class
161+
(
162+
scope == 0 or
163+
not isinstance(
164+
self._scopes[scope - 1].node,
165+
ast.ClassDef,
166+
)
167+
)
168+
):
169+
self.super_offsets.add(ast_to_offset(node))
147170

148171
self.generic_visit(node)
149172

‎tests/features/super_test.py‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,19 @@
2020
'class C(Base):\n'
2121
' def f(self):\n'
2222
' super(Base, self).f()\n',
23+
'class Outer:\n' # common nesting
24+
' class C(Base):\n'
25+
' def f(self):\n'
26+
' super(C, self).f()\n',
27+
'class Outer:\n' # higher levels of nesting
28+
' class Inner:\n'
29+
' class C(Base):\n'
30+
' def f(self):\n'
31+
' super(Inner.C, self).f()\n',
32+
'class Outer:\n' # super arg1 nested in unrelated name
33+
' class C(Base):\n'
34+
' def f(self):\n'
35+
' super(some_module.Outer.C, self).f()\n',
2336
2437
# super outside of a class (technically legal!)
2538
'def f(self):\n'
@@ -87,12 +100,36 @@ def test_fix_super_noop(s):
87100
'class Outer:\n'
88101
' class C(Base):\n'
89102
' def f(self):\n'
90-
' super (C, self).f()\n',
103+
' super (Outer.C, self).f()\n',
91104
'class Outer:\n'
92105
' class C(Base):\n'
93106
' def f(self):\n'
94107
' super().f()\n',
95108
),
109+
(
110+
'def f():\n'
111+
' class Outer:\n'
112+
' class C(Base):\n'
113+
' def f(self):\n'
114+
' super(Outer.C, self).f()\n',
115+
'def f():\n'
116+
' class Outer:\n'
117+
' class C(Base):\n'
118+
' def f(self):\n'
119+
' super().f()\n',
120+
),
121+
(
122+
'class A:\n'
123+
' class B:\n'
124+
' class C:\n'
125+
' def f(self):\n'
126+
' super(A.B.C, self).f()\n',
127+
'class A:\n'
128+
' class B:\n'
129+
' class C:\n'
130+
' def f(self):\n'
131+
' super().f()\n',
132+
),
96133
(
97134
'class C(Base):\n'
98135
' f = lambda self: super(C, self).f()\n',

0 commit comments

Comments
 (0)