Skip to content

Commit

Permalink
checkers.utils.is_sys_guard: understands six.PY2 or six.PY3
Browse files Browse the repository at this point in the history
  • Loading branch information
perrinjerome committed Apr 3, 2024
1 parent 27b1ae7 commit 6e11d88
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 1 deletion.
1 change: 1 addition & 0 deletions doc/whatsnew/fragments/3501.feature
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Understand `six.PY2` and `six.PY3` for conditional imports.
5 changes: 5 additions & 0 deletions pylint/checkers/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1831,6 +1831,11 @@ def is_sys_guard(node: nodes.If) -> bool:
and value.as_string() == "sys.version_info"
):
return True
elif (
isinstance(node.test, nodes.Attribute)
and node.test.as_string() in {'six.PY2', 'six.PY3'}
):
return True

return False

Expand Down
20 changes: 19 additions & 1 deletion tests/checkers/unittest_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -400,9 +400,19 @@ def test_if_sys_guard() -> None:
if sys.some_other_function > (3, 8): #@
pass
import six
if six.PY2: #@
pass
if six.PY3: #@
pass
if six.something_else: #@
pass
"""
)
assert isinstance(code, list) and len(code) == 3
assert isinstance(code, list) and len(code) == 6

assert isinstance(code[0], nodes.If)
assert utils.is_sys_guard(code[0]) is True
Expand All @@ -412,6 +422,14 @@ def test_if_sys_guard() -> None:
assert isinstance(code[2], nodes.If)
assert utils.is_sys_guard(code[2]) is False

assert isinstance(code[3], nodes.If)
assert utils.is_sys_guard(code[3]) is True
assert isinstance(code[4], nodes.If)
assert utils.is_sys_guard(code[4]) is True

assert isinstance(code[5], nodes.If)
assert utils.is_sys_guard(code[5]) is False


def test_if_typing_guard() -> None:
code = astroid.extract_node(
Expand Down

0 comments on commit 6e11d88

Please sign in to comment.