From 31546d2cdf4a7647df46b0e9c4b771efe123ea3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= <cr52@protonmail.com>
Date: Thu, 9 Sep 2021 19:24:12 +0200
Subject: [PATCH] ENH: expand versioned_branches feature to Python 3 minor
 version comparison (<, >, <=, >= with else)

---
 README.md                                 |  41 +++++-
 pyupgrade/_plugins/versioned_branches.py  |  36 ++++-
 tests/features/versioned_branches_test.py | 152 ++++++++++++++++++++++
 3 files changed, 221 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 1366c75d..45f54863 100644
--- a/README.md
+++ b/README.md
@@ -294,21 +294,58 @@ def f():
         yield (a, b)
 ```
 
-### `if PY2` blocks
+### Python2 and old Python3.x blocks
 
 Availability:
 - `--py3-plus` is passed on the commandline.
 
 ```python
 # input
-if six.PY2:      # also understands `six.PY3` and `not` and `sys.version_info`
+import sys
+if sys.version_info < (3,):  # also understands `six.PY2` (and `not`), `six.PY3` (and `not`)
     print('py2')
 else:
     print('py3')
 # output
+import sys
 print('py3')
 ```
 
+Availability:
+- `--py36-plus` will remove Python <= 3.5 only blocks
+- `--py37-plus` will remove Python <= 3.6 only blocks
+- so on and so forth
+
+```python
+# using --py36-plus for this example
+# input
+import sys
+if sys.version_info < (3, 6):
+    print('py3.5')
+else:
+    print('py3.6+')
+
+if sys.version_info <= (3, 5):
+    print('py3.5')
+else:
+    print('py3.6+')
+
+if sys.version_info >= (3, 6):
+    print('py3.6+')
+else:
+    print('py3.5')
+
+# output
+import sys
+print('py3.6+')
+
+print('py3.6+')
+
+print('py3.6+')
+```
+
+Note that `if` blocks without an `else` will not be rewriten as it could introduce a syntax error.
+
 ### remove `six` compatibility code
 
 Availability:
diff --git a/pyupgrade/_plugins/versioned_branches.py b/pyupgrade/_plugins/versioned_branches.py
index df66b744..6cc7a1c4 100644
--- a/pyupgrade/_plugins/versioned_branches.py
+++ b/pyupgrade/_plugins/versioned_branches.py
@@ -14,6 +14,7 @@
 from pyupgrade._data import register
 from pyupgrade._data import State
 from pyupgrade._data import TokenFunc
+from pyupgrade._data import Version
 from pyupgrade._token_helpers import Block
 
 
@@ -77,6 +78,7 @@ 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:
     if not (
             isinstance(test.ops[0], op) and
@@ -87,9 +89,11 @@ def _compare_to_3(
         return False
 
     # checked above but mypy needs help
-    elts = cast('List[ast.Num]', test.comparators[0].elts)
+    ast_elts = cast('List[ast.Num]', test.comparators[0].elts)
+    # padding a 0 for compatibility with (3,) used as a spec
+    elts = tuple(e.n for e in ast_elts) + (0,)
 
-    return elts[0].n == 3 and all(n.n == 0 for n in elts[1:])
+    return elts[:2] == (3, minor) and all(n == 0 for n in elts[2:])
 
 
 @register(ast.If)
@@ -98,8 +102,16 @@ def visit_If(
         node: ast.If,
         parent: ast.AST,
 ) -> Iterable[Tuple[Offset, TokenFunc]]:
+
+    min_version: Version
+    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:
@@ -114,6 +126,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(
@@ -124,7 +137,11 @@ 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])
+                        )
                     )
                 )
             )
@@ -132,7 +149,7 @@ def visit_If(
         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:
@@ -147,6 +164,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(
@@ -157,7 +176,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])
+                        )
                     )
                 )
             )
diff --git a/tests/features/versioned_branches_test.py b/tests/features/versioned_branches_test.py
index b135fd90..3c865591 100644
--- a/tests/features/versioned_branches_test.py
+++ b/tests/features/versioned_branches_test.py
@@ -452,3 +452,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