Skip to content

Commit

Permalink
Fix detection of self argument name when positional-only
Browse files Browse the repository at this point in the history
  • Loading branch information
bwhmather committed Jan 18, 2024
1 parent 3d87560 commit e5890c8
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 3 deletions.
8 changes: 5 additions & 3 deletions src/ssort/_method_requirements.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,12 @@ def get_method_requirements(node: ast.AST) -> Iterable[str]:
def _get_method_requirements_for_function_def(
node: ast.FunctionDef | ast.AsyncFunctionDef,
) -> Iterable[str]:
if not node.args.args:
if node.args.posonlyargs:
self_arg = node.args.posonlyargs[0].arg
elif node.args.args:
self_arg = node.args.args[0].arg
else:
return

self_arg = node.args.args[0].arg

for statement in node.body:
yield from _get_attribute_accesses(statement, self_arg)
10 changes: 10 additions & 0 deletions tests/test_method_requirements.py
Original file line number Diff line number Diff line change
Expand Up @@ -315,3 +315,13 @@ def fun(self):
"""
)
assert reqs == ["a", "b"]


def test_method_requirements_positional_only_self():
reqs = _method_requirements(
"""
def fun(self, /, arg):
return self.method(arg)
"""
)
assert reqs == ["method"]
25 changes: 25 additions & 0 deletions tests/test_ssort.py
Original file line number Diff line number Diff line change
Expand Up @@ -653,3 +653,28 @@ def f():

actual = ssort(original)
assert actual == expected


def test_ssort_self_positional_only():
original = _clean(
"""
class C:
def fun(self, arg, /, notself):
return notself._notdep()
def _notdep(self):
pass
"""
)
expected = _clean(
"""
class C:
def fun(self, arg, /, notself):
return notself._notdep()
def _notdep(self):
pass
"""
)
actual = ssort(original)
assert actual == expected

0 comments on commit e5890c8

Please sign in to comment.