diff --git a/pyupgrade/_plugins/versioned_branches.py b/pyupgrade/_plugins/versioned_branches.py index 05c776ac..1ec4eda6 100644 --- 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 index 6a7e2994..07fcc11b 100644 --- a/tests/features/versioned_branches_test.py +++ b/tests/features/versioned_branches_test.py @@ -607,6 +607,16 @@ 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', + ), ), ) def test_fix_py3x_only_code(s, expected):