Repository: asottile/pyupgrade

Issue:
Rewrite `if PY2:` blocks in `--py3-plus`
Ideally handle these different cases:

```python
# input
from six import PY2
if PY2:
   x = 5
else:
   x = 6
# output
x = 6
```

But also for

```python
sys.version_info[0] == 2
sys.version_info < (3,)
```
etc.

- [x] `six.PY2`
- [x] `six.PY3`
- [x] `not six.PY2`
- [x] `not six.PY3`
- [ ] `sys.version_info >= (3,)`
- [ ] `sys.version_info < (3,)`
- [ ] `sys.version_info[0] == 2`
- [ ] `sys.version_info[0] == 3`

Discussion:

Patch:
diff --git a/README.md b/README.md
index 0420938..11baadc 100644
--- a/README.md
+++ b/README.md
@@ -195,8 +195,8 @@ Availability:
 
 ```python
 def f():
-    for x in y:       # yield from y
-        yield x
+    for x in y:       # yield from x
+        yield y
 
     for a, b in c:    # yield from c
         yield (a, b)
@@ -209,7 +209,7 @@ Availability:
 
 ```python
 # input
-if six.PY2:      # also understands `six.PY3` and `not`
+if six.PY2:      # also understands `six.PY3` and `not` and `sys.version_info`
     print('py2')
 else:
     print('py3')
@@ -217,10 +217,6 @@ else:
 print('py3')
 ```
 
-_note_: this is a work in progress, see [#152]
-
-[#152]: https://github.com/asottile/pyupgrade/issues/152
-
 ### remove `six` compatibility code
 
 Availability:
diff --git a/pyupgrade.py b/pyupgrade.py
index 19d0b61..b83d15e 100644
--- a/pyupgrade.py
+++ b/pyupgrade.py
@@ -1003,6 +1003,7 @@ class FindPy3Plus(ast.NodeVisitor):
 
         self.encode_calls = {}
 
+        self._version_info_imported = False
         self.if_py2_blocks = set()
         self.if_py3_blocks = set()
 
@@ -1036,11 +1037,28 @@ class FindPy3Plus(ast.NodeVisitor):
             node.attr in names
         )
 
+    def _is_version_info(self, node):
+        return (
+            isinstance(node, ast.Name) and
+            node.id == 'version_info' and
+            self._version_info_imported
+        ) or (
+            isinstance(node, ast.Attribute) and
+            isinstance(node.value, ast.Name) and
+            node.value.id == 'sys' and
+            node.attr == 'version_info'
+        )
+
     def visit_ImportFrom(self, node):  # type: (ast.ImportFrom) -> None
         if node.module == 'six':
             for name in node.names:
                 if not name.asname:
                     self._six_from_imports.add(name.name)
+        elif node.module == 'sys' and any(
+            name.name == 'version_info' and not name.asname
+            for name in node.names
+        ):
+            self._version_info_imported = True
         self.generic_visit(node)
 
     def visit_ClassDef(self, node):  # type: (ast.ClassDef) -> None
@@ -1152,6 +1170,25 @@ class FindPy3Plus(ast.NodeVisitor):
 
         self.generic_visit(node)
 
+    @staticmethod
+    def _eq(test, n):  # type: (ast.Compare, int) -> bool
+        return (
+            isinstance(test.ops[0], ast.Eq) and
+            isinstance(test.comparators[0], ast.Num) and
+            test.comparators[0].n == n
+        )
+
+    @staticmethod
+    def _compare_to_3(test, op):
+        return (
+            isinstance(test.ops[0], op) and
+            isinstance(test.comparators[0], ast.Tuple) and
+            len(test.comparators[0].elts) >= 1 and
+            all(isinstance(n, ast.Num) for n in test.comparators[0].elts) and
+            test.comparators[0].elts[0].n == 3 and
+            all(n.n == 0 for n in test.comparators[0].elts[1:])
+        )
+
     def visit_If(self, node):  # type: (ast.If) -> None
         if node.orelse and not isinstance(node.orelse[0], ast.If):
             if (
@@ -1162,6 +1199,15 @@ class FindPy3Plus(ast.NodeVisitor):
                         isinstance(node.test, ast.UnaryOp) and
                         isinstance(node.test.op, ast.Not) and
                         self._is_six(node.test.operand, ('PY3',))
+                    ) or
+                    # sys.version_info == 2 or < (3,)
+                    (
+                        isinstance(node.test, ast.Compare) and
+                        self._is_version_info(node.test.left) and
+                        len(node.test.ops) == 1 and (
+                            self._eq(node.test, 2) or
+                            self._compare_to_3(node.test, ast.Lt)
+                        )
                     )
             ):
                 self.if_py2_blocks.add(_ast_to_offset(node))
@@ -1173,6 +1219,15 @@ class FindPy3Plus(ast.NodeVisitor):
                         isinstance(node.test, ast.UnaryOp) and
                         isinstance(node.test.op, ast.Not) and
                         self._is_six(node.test.operand, ('PY2',))
+                    ) or
+                    # sys.version_info == 3 or >= (3,) or > (3,)
+                    (
+                        isinstance(node.test, ast.Compare) and
+                        self._is_version_info(node.test.left) and
+                        len(node.test.ops) == 1 and (
+                            self._eq(node.test, 3) or
+                            self._compare_to_3(node.test, (ast.Gt, ast.GtE))
+                        )
                     )
             ):
                 self.if_py3_blocks.add(_ast_to_offset(node))


Test patch:
diff --git a/tests/pyupgrade_test.py b/tests/pyupgrade_test.py
index 90b681b..fe3d328 100644
--- a/tests/pyupgrade_test.py
+++ b/tests/pyupgrade_test.py
@@ -1655,6 +1655,17 @@ def test_fix_encode_noop(s):
         '    pass\n'
         'else:\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'
+        '    3\n'
+        'else:\n'
+        '    2\n',
     ),
 )
 def test_fix_py2_block_noop(s):
@@ -1894,6 +1905,78 @@ def test_fix_py2_block_noop(s):
 
             id='six.PY3, comment after',
         ),
+        pytest.param(
+            'if sys.version_info == 2:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            '3\n',
+
+            id='sys.version_info == 2',
+        ),
+        pytest.param(
+            'if sys.version_info < (3,):\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            '3\n',
+
+            id='sys.version_info < (3,)',
+        ),
+        pytest.param(
+            'if sys.version_info < (3, 0):\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            '3\n',
+
+            id='sys.version_info < (3, 0)',
+        ),
+        pytest.param(
+            'if sys.version_info == 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            '3\n',
+
+            id='sys.version_info == 3',
+        ),
+        pytest.param(
+            'if sys.version_info > (3,):\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            '3\n',
+
+            id='sys.version_info > (3,)',
+        ),
+        pytest.param(
+            'if sys.version_info >= (3,):\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            '3\n',
+
+            id='sys.version_info >= (3,)',
+        ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info > (3,):\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, > (3,)',
+        ),
     ),
 )
 def test_fix_py2_blocks(s, expected):