Merge pull request #467 from asottile/elif_indented

fix elif -> else when indented
asottile · Jun 9, 2021 · 7c50135 · 7c50135
2 parents 1328f3d + 9516778
commit 7c50135
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/pyupgrade/_token_helpers.py b/pyupgrade/_token_helpers.py
@@ -218,7 +218,10 @@ def dedent(self, tokens: List[Token]) -> None:
                 tokens[i] = tokens[i]._replace(src=tokens[i].src[diff:])
 
     def replace_condition(self, tokens: List[Token], new: List[Token]) -> None:
-        tokens[self.start:self.colon] = new
+        start = self.start
+        while tokens[start].name == 'UNIMPORTANT_WS':
+            start += 1
+        tokens[start:self.colon] = new
 
     def _trim_end(self, tokens: List[Token]) -> 'Block':
         """the tokenizer reports the end of the block at the beginning of

diff --git a/tests/features/versioned_branches_test.py b/tests/features/versioned_branches_test.py
@@ -415,6 +415,21 @@ def test_fix_py2_block_noop(s):
 
             id='elif six.PY3 no else',
         ),
+        pytest.param(
+            'def f():\n'
+            '    if True:\n'
+            '        print(1)\n'
+            '    elif six.PY3:\n'
+            '        print(3)\n',
+
+            'def f():\n'
+            '    if True:\n'
+            '        print(1)\n'
+            '    else:\n'
+            '        print(3)\n',
+
+            id='elif six.PY3 no else, indented',
+        ),
     ),
 )
 def test_fix_py2_blocks(s, expected):