Resolve #1431: isort incorrectly moving commentsissue/1431/fix-misplaced-comments

3 files changed, 61 insertions, 19 deletions

diff --git a/isort/output.py b/isort/output.py
index 1a6c504b..8cf915f9 100644
--- a/isort/output.py
+++ b/isort/output.py
@@ -328,13 +328,20 @@ def _with_from_imports(
                 while from_imports and from_imports[0] in as_imports:
                     from_import = from_imports.pop(0)
                     as_imports[from_import] = sorting.naturally(as_imports[from_import])
-                    from_comments = parsed.categorized_comments["straight"].get(
-                        f"{module}.{from_import}"
+                    from_comments = (
+                        parsed.categorized_comments["straight"].get(f"{module}.{from_import}") or []
                     )
                     if (
                         parsed.imports[section]["from"][module][from_import]
                         and not only_show_as_imports
                     ):
+                        specific_comment = (
+                            parsed.categorized_comments["nested"]
+                            .get(module, {})
+                            .pop(from_import, None)
+                        )
+                        if specific_comment:
+                            from_comments.append(specific_comment)
                         output.append(
                             wrap.line(
                                 with_comments(
@@ -347,19 +354,31 @@ def _with_from_imports(
                                 config,
                             )
                         )
-                    output.extend(
-                        wrap.line(
-                            with_comments(
-                                from_comments,
-                                import_start + as_import,
-                                removed=config.ignore_comments,
-                                comment_prefix=config.comment_prefix,
-                            ),
-                            parsed.line_separator,
-                            config,
+                        from_comments = []
+
+                    for as_import in as_imports[from_import]:
+                        specific_comment = (
+                            parsed.categorized_comments["nested"]
+                            .get(module, {})
+                            .pop(as_import, None)
                         )
-                        for as_import in as_imports[from_import]
-                    )
+                        if specific_comment:
+                            from_comments.append(specific_comment)
+
+                        output.append(
+                            wrap.line(
+                                with_comments(
+                                    from_comments,
+                                    import_start + as_import,
+                                    removed=config.ignore_comments,
+                                    comment_prefix=config.comment_prefix,
+                                ),
+                                parsed.line_separator,
+                                config,
+                            )
+                        )
+
+                        from_comments = []
 
                 if "*" in from_imports:
                     output.append(
diff --git a/isort/parse.py b/isort/parse.py
index 913cdbb1..4a537422 100644
--- a/isort/parse.py
+++ b/isort/parse.py
@@ -235,7 +235,7 @@ def file_contents(contents: str, config: Config = DEFAULT_CONFIG) -> ParsedConte
                     if (
                         type_of_import == "from"
                         and stripped_line
-                        and " " not in stripped_line
+                        and " " not in stripped_line.replace(" as ", "")
                         and new_comment
                     ):
                         nested_comments[stripped_line] = comments[-1]
@@ -257,7 +257,7 @@ def file_contents(contents: str, config: Config = DEFAULT_CONFIG) -> ParsedConte
                         if (
                             type_of_import == "from"
                             and stripped_line
-                            and " " not in stripped_line
+                            and " " not in stripped_line.replace(" as ", "")
                             and new_comment
                         ):
                             nested_comments[stripped_line] = comments[-1]
@@ -272,7 +272,7 @@ def file_contents(contents: str, config: Config = DEFAULT_CONFIG) -> ParsedConte
                             if (
                                 type_of_import == "from"
                                 and stripped_line
-                                and " " not in stripped_line
+                                and " " not in stripped_line.replace(" as ", "")
                                 and new_comment
                             ):
                                 nested_comments[stripped_line] = comments[-1]
@@ -282,7 +282,7 @@ def file_contents(contents: str, config: Config = DEFAULT_CONFIG) -> ParsedConte
                     if (
                         type_of_import == "from"
                         and stripped_line
-                        and " " not in stripped_line
+                        and " " not in stripped_line.replace(" as ", "")
                         and new_comment
                     ):
                         nested_comments[stripped_line] = comments[-1]
@@ -332,6 +332,15 @@ def file_contents(contents: str, config: Config = DEFAULT_CONFIG) -> ParsedConte
                             pass
                         elif as_name not in as_map["from"][module]:
                             as_map["from"][module].append(as_name)
+
+                        full_name = f"{nested_module} as {as_name}"
+                        associated_comment = nested_comments.get(full_name)
+                        if associated_comment:
+                            categorized_comments["nested"].setdefault(top_level_module, {})[
+                                full_name
+                            ] = associated_comment
+                            if associated_comment in comments:
+                                comments.pop(comments.index(associated_comment))
                     else:
                         module = just_imports[as_index - 1]
                         as_name = just_imports[as_index + 1]
@@ -340,7 +349,7 @@ def file_contents(contents: str, config: Config = DEFAULT_CONFIG) -> ParsedConte
                         elif as_name not in as_map["straight"][module]:
                             as_map["straight"][module].append(as_name)
 
-                    if config.combine_as_imports and nested_module:
+                    if nested_module and config.combine_as_imports:
                         categorized_comments["from"].setdefault(
                             f"{top_level_module}.__combined_as__", []
                         ).extend(comments)
diff --git a/tests/unit/test_regressions.py b/tests/unit/test_regressions.py
index 3b35c401..9a44c967 100644
--- a/tests/unit/test_regressions.py
+++ b/tests/unit/test_regressions.py
@@ -645,3 +645,17 @@ from package import *  # noqa
         force_single_line=True,
         show_diff=True,
     )
+
+
+def test_isort_doesnt_misplace_comments_issue_1431():
+    """Test to ensure isort wont misplace comments.
+    See: https://github.com/PyCQA/isort/issues/1431
+    """
+    input_text = """from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import (
+    MyLovelyCompanyTeamProjectComponent,  # NOT DRY
+)
+from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import (
+    MyLovelyCompanyTeamProjectComponent as component,  # DRY
+)
+"""
+    assert isort.code(input_text, profile="black") == input_text