From 52ea01b4d967bbd9d06fe6d7220c0ccd1e67cd4e Mon Sep 17 00:00:00 2001
From: Timothy Crosley <timothy.crosley@gmail.com>
Date: Sun, 27 Sep 2020 15:06:18 -0700
Subject: [PATCH] Resolved #1502: Improved float-to-top behavior when there is
 an existing import section present at top-of-file.

---
 CHANGELOG.md                         |   1 +
 isort/core.py                        |   3 +
 isort/parse.py                       |   2 +
 tests/unit/test_ticketed_features.py | 122 ++++++++++++++++++++++++---
 4 files changed, 118 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b820a0325..68d114f3a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ Find out more about isort's release policy [here](https://pycqa.github.io/isort/
 ### 5.6.0 TBD
   - Implemented #1433: Provide helpful feedback in case a custom config file is specified without a configuration.
   - Implemented #1494: Default to sorting imports within `.pxd` files.
+  - Implemented #1502: Improved float-to-top behavior when there is an existing import section present at top-of-file.
   - Improved handling of unsupported configuration option errors (see #1475).
   - Fixed #1463: Better interactive documentation for future option.
   - Fixed #1461: Quiet config option not respected by file API in some circumstances.
diff --git a/isort/core.py b/isort/core.py
index d1e945d48..11ae35f95 100644
--- a/isort/core.py
+++ b/isort/core.py
@@ -83,6 +83,9 @@ def process(
                 if line == "# isort: off\n":
                     isort_off = True
                 if current:
+                    if add_imports:
+                        current += line_separator + line_separator.join(add_imports)
+                        add_imports = []
                     parsed = parse.file_contents(current, config=config)
                     extra_space = ""
                     while current and current[-1] == "\n":
diff --git a/isort/parse.py b/isort/parse.py
index d9ab60e2c..6f121c6bd 100644
--- a/isort/parse.py
+++ b/isort/parse.py
@@ -210,6 +210,8 @@ def file_contents(contents: str, config: Config = DEFAULT_CONFIG) -> ParsedConte
             and not lstripped_line.startswith("#")
             and not lstripped_line.startswith("'''")
             and not lstripped_line.startswith('"""')
+            and not lstripped_line.startswith("import")
+            and not lstripped_line.startswith("from")
         ):
             import_index = index - 1
             while import_index and not in_lines[import_index - 1]:
diff --git a/tests/unit/test_ticketed_features.py b/tests/unit/test_ticketed_features.py
index 75b43da89..f952bdc51 100644
--- a/tests/unit/test_ticketed_features.py
+++ b/tests/unit/test_ticketed_features.py
@@ -60,7 +60,8 @@ def my_function_2():
 """,
             float_to_top=True,
         )
-        == """import os
+        == """
+import os
 import sys
 
 
@@ -90,7 +91,8 @@ def my_function_2():
 """,
             float_to_top=True,
         )
-        == """import os
+        == """
+import os
 
 
 def my_function_1():
@@ -105,10 +107,9 @@ def my_function_2():
 """
     )
 
-
-assert (
-    isort.code(
-        """
+    assert (
+        isort.code(
+            """
 import os
 
 
@@ -129,9 +130,10 @@ def my_function_2():
 
 import a
 """,
-        float_to_top=True,
-    )
-    == """import os
+            float_to_top=True,
+        )
+        == """
+import os
 
 
 def my_function_1():
@@ -151,7 +153,7 @@ def y():
 def my_function_2():
     pass
 """
-)
+    )
 
 
 def test_isort_provides_official_api_for_diff_output_issue_1335():
@@ -746,3 +748,103 @@ def test_isort_should_warn_on_empty_custom_config_issue_1433(tmpdir):
     with pytest.warns(None) as warning:
         assert Config(settings_file=str(settings_file)).quiet
     assert not warning
+
+
+def test_float_to_top_should_respect_existing_newlines_between_imports_issue_1502():
+    """When a file has an existing top of file import block before code but after comments
+    isort's float to top feature should respect the existing spacing between the top file comment
+    and the import statements.
+    See: https://github.com/PyCQA/isort/issues/1502
+    """
+    assert isort.check_code(
+        """#!/bin/bash
+'''My comment'''
+
+import a
+
+x = 1
+""",
+        float_to_top=True,
+        show_diff=True,
+    )
+    assert isort.check_code(
+        """#!/bin/bash
+'''My comment'''
+
+
+import a
+
+x = 1
+""",
+        float_to_top=True,
+        show_diff=True,
+    )
+    assert (
+        isort.code(
+            """#!/bin/bash
+'''My comment'''
+
+
+import a
+
+x = 1
+""",
+            float_to_top=True,
+            add_imports=["import b"],
+        )
+        == """#!/bin/bash
+'''My comment'''
+
+
+import a
+import b
+
+x = 1
+"""
+    )
+
+    assert (
+        isort.code(
+            """#!/bin/bash
+'''My comment'''
+
+
+def my_function():
+    pass
+
+
+import a
+""",
+            float_to_top=True,
+        )
+        == """#!/bin/bash
+'''My comment'''
+import a
+
+
+def my_function():
+    pass
+"""
+    )
+
+    assert (
+        isort.code(
+            """#!/bin/bash
+'''My comment'''
+
+
+def my_function():
+    pass
+""",
+            add_imports=["import os"],
+            float_to_top=True,
+        )
+        == """#!/bin/bash
+'''My comment'''
+import os
+
+
+def my_function():
+    pass
+"""
+    )