Repository: asottile/pyupgrade

Issue:
`sys.version_info.major` checks not rewritten
See below example:

```console
$ cat example.py
import sys

# rewritten
if sys.version_info >= (3,):
    print("3")
else:
    print("2")

# not rewritten
if sys.version_info.major >= 3:
    print("3")
else:
    print("2")
```

```console
$ pip freeze | rg pyupgrade
pyupgrade==3.19.1
```

```console
$ pyupgrade --py3-plus example.py                  
Rewriting example.py
```

```console
$ cat example.py 
import sys

# rewritten
print("3")

# not rewritten
if sys.version_info.major >= 3:
    print("3")
else:
    print("2")
```

Real-life example [here](https://github.com/deeplook/sparklines/blob/5e298bf8a99de82d3812ebc25c9267b1a018a750/sparklines/__main__.py#L18).

Discussion:
@akashbajwa: Hey @mxr, can I work on this issue and create PR? :)
@mxr: Sure@asottile: via #1017

Solution patch:
diff --git a/pyupgrade/_plugins/versioned_branches.py b/pyupgrade/_plugins/versioned_branches.py
index 1b6d361..9152bcc 100644
--- a/pyupgrade/_plugins/versioned_branches.py
+++ b/pyupgrade/_plugins/versioned_branches.py
@@ -82,8 +82,20 @@ def _fix_py3_convert_elif(i: int, tokens: list[Token]) -> None:
 
 
 def _eq(test: ast.Compare, n: int) -> bool:
+    return _cmp(test, ast.Eq, n)
+
+
+def _lt(test: ast.Compare, n: int) -> bool:
+    return _cmp(test, ast.Lt, n)
+
+
+def _gte(test: ast.Compare, n: int) -> bool:
+    return _cmp(test, ast.GtE, n)
+
+
+def _cmp(test: ast.Compare, op: type[ast.cmpop], n: int) -> bool:
     return (
-        isinstance(test.ops[0], ast.Eq) and
+        isinstance(test.ops[0], op) and
         isinstance(test.comparators[0], ast.Constant) and
         test.comparators[0].value == n
     )
@@ -164,6 +176,33 @@ def visit_If(
                         for minor in range(min_version[1])
                     )
                 )
+            ) or
+            # sys.version_info[0] == 2 or < 3
+            # sys.version_info.major == 2 or < 3
+            (
+                isinstance(node.test, ast.Compare) and
+                (
+                    (
+                        isinstance(node.test.left, ast.Subscript) and
+                        isinstance(node.test.left.slice, ast.Constant) and
+                        node.test.left.slice.value == 0
+                    ) or
+                    (
+                        isinstance(node.test.left, ast.Attribute) and
+                        node.test.left.attr == 'major'
+                    )
+                ) and
+                is_name_attr(
+                    node.test.left.value,
+                    state.from_imports,
+                    ('sys',),
+                    ('version_info',),
+                ) and
+                len(node.test.ops) == 1 and
+                (
+                    _eq(node.test, 2) or
+                    _lt(node.test, 3)
+                )
             )
     ):
         if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If):
@@ -211,6 +250,33 @@ def visit_If(
                         for minor in range(min_version[1])
                     )
                 )
+            ) or
+            # sys.version_info[0] == 3 or >= 3
+            # sys.version_info.major == 3 or >= 3
+            (
+                isinstance(node.test, ast.Compare) and
+                (
+                    (
+                        isinstance(node.test.left, ast.Subscript) and
+                        isinstance(node.test.left.slice, ast.Constant) and
+                        node.test.left.slice.value == 0
+                    ) or
+                    (
+                        isinstance(node.test.left, ast.Attribute) and
+                        node.test.left.attr == 'major'
+                    )
+                ) and
+                is_name_attr(
+                    node.test.left.value,
+                    state.from_imports,
+                    ('sys',),
+                    ('version_info',),
+                ) and
+                len(node.test.ops) == 1 and
+                (
+                    _eq(node.test, 3) or
+                    _gte(node.test, 3)
+                )
             )
     ):
         if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If):


Test patch:
diff --git a/tests/features/versioned_branches_test.py b/tests/features/versioned_branches_test.py
index 6a7e299..bc15cc9 100644
--- a/tests/features/versioned_branches_test.py
+++ b/tests/features/versioned_branches_test.py
@@ -316,6 +316,46 @@ def test_fix_py2_block_noop(s):
 
             id='sys.version_info == 2',
         ),
+        pytest.param(
+            'if sys.version_info[0] == 2:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            '3\n',
+
+            id='sys.version_info[0] == 2',
+        ),
+        pytest.param(
+            'if sys.version_info[0] < 3:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            '3\n',
+
+            id='sys.version_info[0] < 3',
+        ),
+        pytest.param(
+            'if sys.version_info.major == 2:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            '3\n',
+
+            id='sys.version_info.major == 2',
+        ),
+        pytest.param(
+            'if sys.version_info.major < 3:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            '3\n',
+
+            id='sys.version_info.major < 3',
+        ),
         pytest.param(
             'if sys.version_info < (3,):\n'
             '    2\n'
@@ -366,6 +406,46 @@ def test_fix_py2_block_noop(s):
 
             id='sys.version_info >= (3,)',
         ),
+        pytest.param(
+            'if sys.version_info[0] == 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            '3\n',
+
+            id='sys.version_info[0]== 3',
+        ),
+        pytest.param(
+            'if sys.version_info[0] >= 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            '3\n',
+
+            id='sys.version_info[0] >= 3',
+        ),
+        pytest.param(
+            'if sys.version_info.major == 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            '3\n',
+
+            id='testme sys.version_info.major == 3',
+        ),
+        pytest.param(
+            'if sys.version_info.major >= 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            '3\n',
+
+            id='sys.version_info.major >= 3',
+        ),
         pytest.param(
             'from sys import version_info\n'
             'if version_info > (3,):\n'
@@ -378,6 +458,102 @@ def test_fix_py2_block_noop(s):
 
             id='from sys import version_info, > (3,)',
         ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info[0] == 2:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, [0] == 2',
+        ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info[0] < 3:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, [0] < 3',
+        ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info.major == 2:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, .major == 2',
+        ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info.major < 3:\n'
+            '    2\n'
+            'else:\n'
+            '    3\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, .major < 3',
+        ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info[0] == 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, [0] == 3',
+        ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info[0] >= 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, [0] >= 3',
+        ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info.major == 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, .major == 3',
+        ),
+        pytest.param(
+            'from sys import version_info\n'
+            'if version_info.major >= 3:\n'
+            '    3\n'
+            'else:\n'
+            '    2\n',
+
+            'from sys import version_info\n'
+            '3\n',
+
+            id='from sys import version_info, .major >= 3',
+        ),
         pytest.param(
             'if True:\n'
             '    print(1)\n'