diff --git a/isort/core.py b/isort/core.py index 46bd77f1..b87df676 100644 --- a/isort/core.py +++ b/isort/core.py @@ -68,6 +68,7 @@ def process( first_comment_index_end: int = -1 contains_imports: bool = False in_top_comment: bool = False + top_comment_started: bool = False first_import_section: bool = True indent: str = "" isort_off: bool = False @@ -195,12 +196,23 @@ def process( ] if ( - (index == 0 or (index in {1, 2} and not contains_imports)) + ( + index == 0 + or ( + index in {1, 2} + and not contains_imports + and ( + top_comment_started + or stripped_line.startswith("# isort:") + ) + ) + ) and stripped_line.startswith("#") and stripped_line not in config.section_comments and stripped_line not in CODE_SORT_COMMENTS ): in_top_comment = True + top_comment_started = True elif in_top_comment and ( not line.startswith("#") or stripped_line in config.section_comments @@ -303,6 +315,8 @@ def process( and stripped_line not in config.treat_comments_as_code ): import_section += line + if not indent and stripped_line and index < 3 and not top_comment_started: + indent = line[: -len(line.lstrip())] elif stripped_line.startswith(IMPORT_START_IDENTIFIERS): new_indent = line[: -len(line.lstrip())] import_statement = line diff --git a/tests/unit/test_regressions.py b/tests/unit/test_regressions.py index ce959ac5..5f7eda41 100644 --- a/tests/unit/test_regressions.py +++ b/tests/unit/test_regressions.py @@ -2094,3 +2094,55 @@ def test_noqa_wrap_mode_idempotent_with_existing_comment(): second_pass = isort.code(first_pass, multi_line_output=7) assert second_pass == first_pass assert second_pass.count("NOQA") == 1 + + +def test_pylint_disable_next_stays_with_import_issue_2054(): + """pylint disable-next comments should stay associated with their imports after sorting. + + When a function body begins with ``# pylint: disable-next=...`` at the first or second + line position, isort was incorrectly treating it as a file-level top comment and writing + it to the output before sorting the imports. Each comment must travel with the specific + import that follows it. + + See: https://github.com/PyCQA/isort/issues/2054 + """ + test_input = """def foo(): + # pylint: disable-next=no-name-in-module + from C import D + # pylint: disable-next=no-name-in-module + from A import B +""" + assert isort.code(test_input) == """def foo(): + # pylint: disable-next=no-name-in-module + from A import B + # pylint: disable-next=no-name-in-module + from C import D +""" + + # Single comment with one import that sorts after another import + test_input_single = """def bar(): + # pylint: disable-next=no-name-in-module + from Z import W + import A +""" + assert isort.code(test_input_single) == """def bar(): + import A + # pylint: disable-next=no-name-in-module + from Z import W +""" + + # Comment preceded by a blank line inside the function body + test_input_blank = """def baz(): + + # pylint: disable-next=no-name-in-module + from C import D + from A import B +""" + # After sorting, comment stays with the import it annotated (from C import D), + # while the unannotated import (from A import B) sorts before it. + assert isort.code(test_input_blank) == """def baz(): + + from A import B + # pylint: disable-next=no-name-in-module + from C import D +"""