Simplify core sorting logic

1 files changed, 82 insertions, 97 deletions

diff --git a/isort/api.py b/isort/api.py
index e78b42c4..4c6bc949 100644
--- a/isort/api.py
+++ b/isort/api.py
@@ -1,7 +1,8 @@
 import re
 from io import StringIO
+from itertools import chain
 from pathlib import Path
-from typing import Any, NamedTuple, Optional, TextIO, Tuple
+from typing import Any, List, NamedTuple, Optional, TextIO, Tuple
 
 from . import output, parse
 from .exceptions import (
@@ -144,7 +145,7 @@ def sort_imports(
     - `config`: Config settings to use when sorting imports. Defaults settings.DEFAULT_CONFIG.
     """
     line_separator: str = config.line_ending
-    add_imports = [format_natural(addition) for addition in config.add_imports]
+    add_imports: List[str] = [format_natural(addition) for addition in config.add_imports]
     import_section: str = ""
     in_quote: str = ""
     first_comment_index_start: int = -1
@@ -154,98 +155,99 @@ def sort_imports(
     first_import_section: bool = True
     section_comments = [f"# {heading}" for heading in config.import_headings.values()]
 
-    def additional_imports() -> str:
-        nonlocal add_imports
-        nonlocal line_separator
-        nonlocal contains_imports
+    for index, line in enumerate(chain(input_stream, (None,))):
+        if line is None:
+            if index == 0 and not config.force_adds:
+                return
 
-        if not add_imports:
-            return ""
-
-        if not line_separator:
-            line_separator = "\n"
-
-        fomatted_imports: str = line_separator.join(add_imports) + line_separator
-        contains_imports = True
-        add_imports = []
-        return fomatted_imports
-
-    index: int = 0
-    for index, line in enumerate(input_stream):
-        if not line_separator:
-            line_separator = line[-1]
-
-        stripped_line = line.strip()
-        if (
-            (index == 0 or (index == 1 and not contains_imports))
-            and line.startswith("#")
-            and stripped_line not in section_comments
-        ):
-            in_top_comment = True
-        elif in_top_comment:
-            if not line.startswith("#") or stripped_line in section_comments:
-                in_top_comment = False
-                first_comment_index_end = index - 1
-
-        if not line.startswith("#") and '"' in line or "'" in line:
-            char_index = 0
-            if first_comment_index_start == -1 and (line.startswith('"') or line.startswith("'")):
-                first_comment_index_start = index
-            while char_index < len(line):
-                if line[char_index] == "\\":
-                    char_index += 1
-                elif in_quote:
-                    if line[char_index : char_index + len(in_quote)] == in_quote:
-                        in_quote = ""
-                        if first_comment_index_end < first_comment_index_start:
-                            first_comment_index_end = index
-                elif line[char_index] in ("'", '"'):
-                    long_quote = line[char_index : char_index + 3]
-                    if long_quote in ('"""', "'''"):
-                        in_quote = long_quote
-                        char_index += 2
-                    else:
-                        in_quote = line[char_index]
-                elif line[char_index] == "#":
-                    break
-                char_index += 1
+            not_imports = True
+            line = ""
+            if not line_separator:
+                line_separator = "\n"
+        else:
+            if not line_separator:
+                line_separator = line[-1]
 
-        not_imports = bool(in_quote) or in_top_comment
-        if not (in_quote or in_top_comment):
             stripped_line = line.strip()
-            if not stripped_line or stripped_line.startswith("#"):
-                import_section += line
-            elif stripped_line.startswith(IMPORT_START_IDENTIFIERS):
-                import_section += line
-                while stripped_line.endswith("\\") or (
-                    "(" in stripped_line and ")" not in stripped_line
+            if (
+                (index == 0 or (index == 1 and not contains_imports))
+                and line.startswith("#")
+                and stripped_line not in section_comments
+            ):
+                in_top_comment = True
+            elif in_top_comment:
+                if not line.startswith("#") or stripped_line in section_comments:
+                    in_top_comment = False
+                    first_comment_index_end = index - 1
+
+            if not line.startswith("#") and '"' in line or "'" in line:
+                char_index = 0
+                if first_comment_index_start == -1 and (
+                    line.startswith('"') or line.startswith("'")
                 ):
-                    if stripped_line.endswith("\\"):
-                        while stripped_line and stripped_line.endswith("\\"):
-                            line = input_stream.readline()
-                            stripped_line = line.strip().split("#")[0]
-                            import_section += line
-                    else:
-                        while ")" not in stripped_line:
-                            line = input_stream.readline()
-                            stripped_line = line.strip().split("#")[0]
-                            import_section += line
+                    first_comment_index_start = index
+                while char_index < len(line):
+                    if line[char_index] == "\\":
+                        char_index += 1
+                    elif in_quote:
+                        if line[char_index : char_index + len(in_quote)] == in_quote:
+                            in_quote = ""
+                            if first_comment_index_end < first_comment_index_start:
+                                first_comment_index_end = index
+                    elif line[char_index] in ("'", '"'):
+                        long_quote = line[char_index : char_index + 3]
+                        if long_quote in ('"""', "'''"):
+                            in_quote = long_quote
+                            char_index += 2
+                        else:
+                            in_quote = line[char_index]
+                    elif line[char_index] == "#":
+                        break
+                    char_index += 1
 
-                contains_imports = True
-            else:
-                not_imports = True
+            not_imports = bool(in_quote) or in_top_comment
+            if not (in_quote or in_top_comment):
+                stripped_line = line.strip()
+                if not stripped_line or stripped_line.startswith("#"):
+                    import_section += line
+                elif stripped_line.startswith(IMPORT_START_IDENTIFIERS):
+                    import_section += line
+                    while stripped_line.endswith("\\") or (
+                        "(" in stripped_line and ")" not in stripped_line
+                    ):
+                        if stripped_line.endswith("\\"):
+                            while stripped_line and stripped_line.endswith("\\"):
+                                line = input_stream.readline()
+                                stripped_line = line.strip().split("#")[0]
+                                import_section += line
+                        else:
+                            while ")" not in stripped_line:
+                                line = input_stream.readline()
+                                stripped_line = line.strip().split("#")[0]
+                                import_section += line
+
+                    contains_imports = True
+                else:
+                    not_imports = True
 
         if not_imports:
             if (
-                not in_top_comment
+                add_imports
+                and not in_top_comment
                 and not in_quote
                 and not import_section
                 and not line.lstrip().startswith(COMMENT_INDICATORS)
             ):
-                import_section = additional_imports()
+                import_section = line_separator.join(add_imports) + line_separator
+                contains_imports = True
+                add_imports = []
 
             if import_section:
-                import_section += additional_imports()
+                if add_imports:
+                    import_section += line_separator.join(add_imports) + line_separator
+                    contains_imports = True
+                    add_imports = []
+
                 import_section += line
                 if not contains_imports:
                     output_stream.write(import_section)
@@ -265,20 +267,3 @@ def sort_imports(
             else:
                 output_stream.write(line)
                 not_imports = False
-
-    if import_section:
-        import_section += additional_imports()
-        if not contains_imports:
-            output_stream.write(import_section)
-        else:
-            if first_import_section and not import_section.lstrip(line_separator).startswith(
-                COMMENT_INDICATORS
-            ):
-                import_section = import_section.lstrip(line_separator)
-            output_stream.write(
-                output.sorted_imports(
-                    parse.file_contents(import_section, config=config), config, extension
-                )
-            )
-    elif index > 1 or config.force_adds:
-        output_stream.write(additional_imports())