Merge pull request #858 from asottile/versioned-branches-elif

handle more flavors of versioned branches

pyupgrade/_plugins/versioned_branches.py

@@ -45,6 +45,20 @@ def _fix_py2_block(i: int, tokens: list[Token]) -> None:
         del tokens[if_block.start:else_block.start]
 
 
+def _fix_remove_block(i: int, tokens: list[Token]) -> None:
+    block = Block.find(tokens, i)
+    del tokens[block.start:block.end]
+
+
+def _fix_py2_convert_elif(i: int, tokens: list[Token]) -> None:
+    if_block = Block.find(tokens, i)
+    # wasn't actually followed by an `elif`
+    if tokens[if_block.end].src != 'elif':
+        return
+    tokens[if_block.end] = Token('CODE', tokens[i].src)
+    _fix_remove_block(i, tokens)
+
+
 def _fix_py3_block_else(i: int, tokens: list[Token]) -> None:
     if tokens[i].src == 'if':
         if_block, else_block = _find_if_else_block(tokens, i)
@@ -57,9 +71,14 @@ def _fix_py3_block_else(i: int, tokens: list[Token]) -> None:
         if_block.replace_condition(tokens, [Token('NAME', 'else')])
 
 
-def _fix_remove_block(i: int, tokens: list[Token]) -> None:
-    block = Block.find(tokens, i)
-    del tokens[block.start:block.end]
+def _fix_py3_convert_elif(i: int, tokens: list[Token]) -> None:
+    if_block = Block.find(tokens, i)
+    # wasn't actually followed by an `elif`
+    if tokens[if_block.end].src != 'elif':
+        return
+    tokens[if_block.end] = Token('CODE', tokens[i].src)
+    if_block.dedent(tokens)
+    del tokens[if_block.start:if_block.block]
 
 
 def _eq(test: ast.Compare, n: int) -> bool:
@@ -147,9 +166,11 @@ def visit_If(
                 )
             )
     ):
-        if node.orelse and not isinstance(node.orelse[0], ast.If):
+        if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If):
+            yield ast_to_offset(node), _fix_py2_convert_elif
+        elif node.orelse:
             yield ast_to_offset(node), _fix_py2_block
-        elif node.col_offset == 0 and not node.orelse:
+        elif node.col_offset == 0:
             yield ast_to_offset(node), _fix_remove_block
     elif (
             # if six.PY3:
@@ -192,7 +213,9 @@ def visit_If(
                 )
             )
     ):
-        if node.orelse and not isinstance(node.orelse[0], ast.If):
+        if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If):
+            yield ast_to_offset(node), _fix_py3_convert_elif
+        elif node.orelse:
             yield ast_to_offset(node), _fix_py3_block_else
-        elif not node.orelse:
+        else:
             yield ast_to_offset(node), _fix_py3_block

tests/features/versioned_branches_test.py

@@ -13,15 +13,22 @@
         'if True:\n'
         '    if six.PY2:\n'
         '        pass\n',
-        # for now we don't attempt to rewrite `elif`
-        'if six.PY2:\n'
-        '    pass\n'
-        'elif False:\n'
-        '    pass\n',
-        'if six.PY3:\n'
-        '    pass\n'
-        'elif False:\n'
-        '    pass\n',
+        pytest.param(
+            'if six.PY2:\n'
+            '    2\n'
+            'else:\n'
+            '    if False:\n'
+            '        3\n',
+            id='py2 indistinguisable at ast level from `elif`',
+        ),
+        pytest.param(
+            'if six.PY3:\n'
+            '    3\n'
+            'else:\n'
+            '    if False:\n'
+            '        2\n',
+            id='py3 indistinguisable at ast level from `elif`',
+        ),
         # don't rewrite version compares with not 3.0 compares
         'if sys.version_info >= (3, 6):\n'
         '    3.6\n'
@@ -443,6 +450,27 @@ def test_fix_py2_block_noop(s):
 
             id='comment after dedented block',
         ),
+        pytest.param(
+            'print("before")\n'
+            'if six.PY2:\n'
+            '    pass\n'
+            'print("after")\n',
+
+            'print("before")\n'
+            'print("after")\n',
+            id='can remove no-else if at module scope',
+        ),
+        pytest.param(
+            'if six.PY2:\n'
+            '    pass\n'
+            'elif False:\n'
+            '    pass\n',
+
+            'if False:\n'
+            '    pass\n',
+
+            id='elif becomes if',
+        ),
     ),
 )
 def test_fix_py2_blocks(s, expected):
@@ -568,14 +596,16 @@ def test_fix_py3_only_code(s, expected):
             id='sys.version_info >= (3, 6), noelse',
         ),
         pytest.param(
-            'print("before")\n'
-            'if six.PY2:\n'
+            'if six.PY3:\n'
             '    pass\n'
-            'print("after")\n',
+            'elif False:\n'
+            '    pass\n',
 
-            'print("before")\n'
-            'print("after")\n',
-            id='can remove no-else if at module scope',
+            'pass\n'
+            'if False:\n'
+            '    pass\n',
+
+            id='elif becomes if',
         ),
     ),
 )