Fix pylint disable-next comments detaching from imports at function-body start

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

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
+"""