Merge pull request #362 from asottile/invalid_indent_comment_after_block

fix invalid dedent with comment after block
asottile · Sep 25, 2021 · 45b4109 · 45b4109
2 parents 9c40758 + 2fce4a1
commit 45b4109
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/pyupgrade/_token_helpers.py b/pyupgrade/_token_helpers.py
@@ -195,7 +195,9 @@ def _minimum_indent(self, tokens: List[Token]) -> int:
         for i in range(self.block, self.end):
             if (
                     tokens[i - 1].name in ('NL', 'NEWLINE') and
-                    tokens[i].name in ('INDENT', UNIMPORTANT_WS)
+                    tokens[i].name in ('INDENT', UNIMPORTANT_WS) and
+                    # comments can have arbitrary indentation so ignore them
+                    tokens[i + 1].name != 'COMMENT'
             ):
                 token_indent = len(tokens[i].src)
                 if block_indent is None:
@@ -209,13 +211,17 @@ def _minimum_indent(self, tokens: List[Token]) -> int:
     def dedent(self, tokens: List[Token]) -> None:
         if self.line:
             return
-        diff = self._minimum_indent(tokens) - self._initial_indent(tokens)
+        initial_indent = self._initial_indent(tokens)
+        diff = self._minimum_indent(tokens) - initial_indent
         for i in range(self.block, self.end):
             if (
                     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:])
+                # make sure we preserve *at least* the initial indent
+                s = tokens[i].src
+                s = s[:initial_indent] + s[initial_indent + diff:]
+                tokens[i] = tokens[i]._replace(src=s)
 
     def replace_condition(self, tokens: List[Token], new: List[Token]) -> None:
         start = self.start

diff --git a/tests/features/versioned_branches_test.py b/tests/features/versioned_branches_test.py
@@ -430,6 +430,20 @@ def test_fix_py2_block_noop(s):
 
             id='elif six.PY3 no else, indented',
         ),
+        pytest.param(
+            'if True:\n'
+            '    if sys.version_info > (3,):\n'
+            '        print(3)\n'
+            '    # comment\n'
+            '    print(2+3)\n',
+
+            'if True:\n'
+            '    print(3)\n'
+            '    # comment\n'
+            '    print(2+3)\n',
+
+            id='comment after dedented block',
+        ),
     ),
 )
 def test_fix_py2_blocks(s, expected):