Fix removal of comments after dropped version block

asottile · Oct 10, 2024 · 68e6854 · 68e6854
1 parent 99f0d7a
commit 68e6854
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/pyupgrade/_plugins/versioned_branches.py b/pyupgrade/_plugins/versioned_branches.py
@@ -46,7 +46,7 @@ def _fix_py2_block(i: int, tokens: list[Token]) -> None:
 
 
 def _fix_remove_block(i: int, tokens: list[Token]) -> None:
-    block = Block.find(tokens, i)
+    block = Block.find(tokens, i, trim_end=True)
     del tokens[block.start:block.end]
 
 

diff --git a/tests/features/versioned_branches_test.py b/tests/features/versioned_branches_test.py
@@ -607,6 +607,27 @@ def test_fix_py3_only_code(s, expected):
 
             id='elif becomes if',
         ),
+        pytest.param(
+            'import sys\n'
+            'if sys.version_info < (3, 6):\n'
+            '    3-5\n'
+            '# comment',
+
+            'import sys\n'
+            '# comment',
+            id='sys.version_info < (3, 6), trailing comment',
+        ),
+        pytest.param(
+            'import sys\n'
+            'if sys.version_info < (3, 6):\n'
+            '    3-5\n'
+            '    # comment here\n'
+            '# comment there',
+
+            'import sys\n'
+            '# comment there',
+            id='sys.version_info < (3, 6), two comments',
+        ),
     ),
 )
 def test_fix_py3x_only_code(s, expected):