Skip to content

Commit

Permalink
ENH: expand versioned_branches feature to Python 3 minor version comp…
Browse files Browse the repository at this point in the history 
…arison (<, >, <=, >= with else)
  • Loading branch information
@neutrinoceros
neutrinoceros committed Sep 11, 2021
1 parent f86dc65 commit c2696c2
Show file tree
Hide file tree
Showing 2 changed files with 186 additions and 11 deletions.
40 changes: 34 additions & 6 deletions pyupgrade/_plugins/versioned_branches.py
View file
Expand Up @@ -77,19 +77,27 @@ def _eq(test: ast.Compare, n: int) -> bool:
def _compare_to_3(
test: ast.Compare,
op: Union[Type[ast.cmpop], Tuple[Type[ast.cmpop], ...]],
minor: int = 0,
) -> bool:
min_len = 2 if minor else 1
if not (
isinstance(test.ops[0], op) and
isinstance(test.comparators[0], ast.Tuple) and
len(test.comparators[0].elts) >= 1 and
len(test.comparators[0].elts) >= min_len and
all(isinstance(n, ast.Num) for n in test.comparators[0].elts)
):
return False

# checked above but mypy needs help
elts = cast('List[ast.Num]', test.comparators[0].elts)

return elts[0].n == 3 and all(n.n == 0 for n in elts[1:])
retv = elts[0].n == 3
offset = 1
if minor:
retv &= elts[1].n == minor
offset += 1
retv &= all(n.n == 0 for n in elts[offset:])
return retv


@register(ast.If)
Expand All @@ -98,8 +106,16 @@ def visit_If(
node: ast.If,
parent: ast.AST,
) -> Iterable[Tuple[Offset, TokenFunc]]:

min_version: Tuple[int, ...]
if state.settings.min_version == (3,):
min_version = (3, 0)
else:
min_version = state.settings.min_version
assert len(min_version) >= 2

if (
state.settings.min_version >= (3,) and (
min_version >= (3,) and (
# if six.PY2:
is_name_attr(node.test, state.from_imports, 'six', ('PY2',)) or
# if not six.PY3:
Expand All @@ -114,6 +130,7 @@ def visit_If(
)
) or
# sys.version_info == 2 or < (3,)
# or < (3, n) or <= (3, n) (with n<m)
(
isinstance(node.test, ast.Compare) and
is_name_attr(
Expand All @@ -124,15 +141,19 @@ def visit_If(
) and
len(node.test.ops) == 1 and (
_eq(node.test, 2) or
_compare_to_3(node.test, ast.Lt)
_compare_to_3(node.test, ast.Lt, min_version[1]) or
any(
_compare_to_3(node.test, (ast.Lt, ast.LtE), minor)
for minor in range(min_version[1])
)
)
)
)
):
if node.orelse and not isinstance(node.orelse[0], ast.If):
yield ast_to_offset(node), _fix_py2_block
elif (
state.settings.min_version >= (3,) and (
min_version >= (3,) and (
# if six.PY3:
is_name_attr(node.test, state.from_imports, 'six', ('PY3',)) or
# if not six.PY2:
Expand All @@ -147,6 +168,8 @@ def visit_If(
)
) or
# sys.version_info == 3 or >= (3,) or > (3,)
# sys.version_info >= (3, n) (with n<=m)
# or sys.version_info > (3, n) (with n<m)
(
isinstance(node.test, ast.Compare) and
is_name_attr(
Expand All @@ -157,7 +180,12 @@ def visit_If(
) and
len(node.test.ops) == 1 and (
_eq(node.test, 3) or
_compare_to_3(node.test, (ast.Gt, ast.GtE))
_compare_to_3(node.test, (ast.Gt, ast.GtE)) or
_compare_to_3(node.test, ast.GtE, min_version[1]) or
any(
_compare_to_3(node.test, (ast.Gt, ast.GtE), minor)
for minor in range(min_version[1])
)
)
)
)
Expand Down
157 changes: 152 additions & 5 deletions tests/features/versioned_branches_test.py
View file
Expand Up @@ -23,11 +23,6 @@
' pass\n'
'elif False:\n'
' pass\n',
# don't rewrite version compares with not 3.0 compares
'if sys.version_info >= (3, 6):\n'
' 3.6\n'
'else:\n'
' 3.5\n',
# don't try and think about `sys.version`
'from sys import version\n'
'if sys.version[0] > "2":\n'
Expand Down Expand Up @@ -452,3 +447,155 @@ def test_fix_py2_blocks(s, expected):
def test_fix_py3_only_code(s, expected):
ret = _fix_plugins(s, settings=Settings(min_version=(3,)))
assert ret == expected


@pytest.mark.parametrize(
('s', 'expected'),
(
pytest.param(
'import sys\n'
'if sys.version_info > (3, 5):\n'
' 3+6\n'
'else:\n'
' 3-5\n',
'import sys\n'
'3+6\n',
id='sys.version_info > (3, 5)',
),
pytest.param(
'from sys import version_info\n'
'if version_info > (3, 5):\n'
' 3+6\n'
'else:\n'
' 3-5\n',
'from sys import version_info\n'
'3+6\n',
id='from sys import version_info, > (3, 5)',
),
pytest.param(
'import sys\n'
'if sys.version_info >= (3, 6):\n'
' 3+6\n'
'else:\n'
' 3-5\n',
'import sys\n'
'3+6\n',
id='sys.version_info >= (3, 6)',
),
pytest.param(
'from sys import version_info\n'
'if version_info >= (3, 6):\n'
' 3+6\n'
'else:\n'
' 3-5\n',
'from sys import version_info\n'
'3+6\n',
id='from sys import version_info, >= (3, 6)',
),
pytest.param(
'import sys\n'
'if sys.version_info < (3, 6):\n'
' 3-5\n'
'else:\n'
' 3+6\n',
'import sys\n'
'3+6\n',
id='sys.version_info < (3, 6)',
),
pytest.param(
'from sys import version_info\n'
'if version_info < (3, 6):\n'
' 3-5\n'
'else:\n'
' 3+6\n',
'from sys import version_info\n'
'3+6\n',
id='from sys import version_info, < (3, 6)',
),
pytest.param(
'import sys\n'
'if sys.version_info <= (3, 5):\n'
' 3-5\n'
'else:\n'
' 3+6\n',
'import sys\n'
'3+6\n',
id='sys.version_info <= (3, 5)',
),
pytest.param(
'from sys import version_info\n'
'if version_info <= (3, 5):\n'
' 3-5\n'
'else:\n'
' 3+6\n',
'from sys import version_info\n'
'3+6\n',
id='from sys import version_info, <= (3, 5)',
),
),
)
def test_fix_py3x_only_code(s, expected):
ret = _fix_plugins(s, settings=Settings(min_version=(3, 6)))
assert ret == expected


@pytest.mark.parametrize(
's',
(
# we timidly skip `if` without `else` as it could cause a SyntaxError
'import sys'
'if sys.version_info >= (3, 6):\n'
' pass',
# here's the case where it causes a SyntaxError
'import sys'
'if True'
' if sys.version_info >= (3, 6):\n'
' pass\n',
# both branches are still relevant in the following cases
'import sys\n'
'if sys.version_info > (3, 7):\n'
' 3-6\n'
'else:\n'
' 3+7\n',
'import sys\n'
'if sys.version_info < (3, 7):\n'
' 3-6\n'
'else:\n'
' 3+7\n',
'import sys\n'
'if sys.version_info >= (3, 7):\n'
' 3+7\n'
'else:\n'
' 3-6\n',
'import sys\n'
'if sys.version_info <= (3, 7):\n'
' 3-7\n'
'else:\n'
' 3+8\n',
'import sys\n'
'if sys.version_info <= (3, 6):\n'
' 3-6\n'
'else:\n'
' 3+7\n',
'import sys\n'
'if sys.version_info > (3, 6):\n'
' 3+7\n'
'else:\n'
' 3-6\n',
),
)
def test_fix_py3x_only_noop(s):
assert _fix_plugins(s, settings=Settings(min_version=(3, 6))) == s

0 comments on commit c2696c2

Please sign in to comment.