Merge pull request #319 from asottile/fix_py2_branches

fix py2 branches with multiple internal blocks
asottile · Jun 12, 2020 · c177b69 · c177b69
2 parents 4656fbe + 7290af2
commit c177b69
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/pyupgrade.py b/pyupgrade.py
@@ -1771,7 +1771,7 @@ def dedent(self, tokens: List[Token]) -> None:
         diff = self._minimum_indent(tokens) - self._initial_indent(tokens)
         for i in range(self.block, self.end):
             if (
-                    tokens[i - 1].name in ('NL', 'NEWLINE') and
+                    tokens[i - 1].name in ('DEDENT', 'NL', 'NEWLINE') and
                     tokens[i].name in ('INDENT', UNIMPORTANT_WS)
             ):
                 tokens[i] = tokens[i]._replace(src=tokens[i].src[diff:])

diff --git a/tests/versioned_branches_test.py b/tests/versioned_branches_test.py
@@ -264,6 +264,25 @@ def test_fix_py2_block_noop(s):
 
             id='six.PY2, comment after',
         ),
+        pytest.param(
+            'if six.PY2:\n'
+            '    def f():\n'
+            '        print("py2")\n'
+            '    def g():\n'
+            '        print("py2")\n'
+            'else:\n'
+            '    def f():\n'
+            '        print("py3")\n'
+            '    def g():\n'
+            '        print("py3")\n',
+
+            'def f():\n'
+            '    print("py3")\n'
+            'def g():\n'
+            '    print("py3")\n',
+
+            id='six.PY2 multiple functions',
+        ),
         pytest.param(
             'if True:\n'
             '    if six.PY3:\n'