Preserve blank lines between pre-import code and first import

reorder_python_imports.py

@@ -107,9 +107,16 @@ def partition_source(src: str) -> tuple[str, list[str], str, str]:
         else:
             chunks.append((CodeType.CODE, s))
 
+    last_pre_import_idx = -1
+    first_import_idx = None
     last_idx = 0
     for i, (tp, _) in enumerate(chunks):
-        if tp in (CodeType.PRE_IMPORT_CODE, CodeType.IMPORT):
+        if tp is CodeType.PRE_IMPORT_CODE:
+            last_pre_import_idx = i
+            last_idx = i
+        elif tp is CodeType.IMPORT:
+            if first_import_idx is None:
+                first_import_idx = i
             last_idx = i
 
     pre = []
@@ -118,6 +125,13 @@ def partition_source(src: str) -> tuple[str, list[str], str, str]:
     for i, (tp, src) in enumerate(chunks):
         if tp is CodeType.PRE_IMPORT_CODE:
             pre.append(src)
+        elif (
+                tp is CodeType.NON_CODE and
+                first_import_idx is not None and
+                i > last_pre_import_idx and
+                i < first_import_idx
+        ):
+            pre.append(src)
         elif tp is CodeType.IMPORT:
             imports.append(src)
         elif tp is CodeType.CODE or i > last_idx:

tests/reorder_python_imports_test.py

@@ -205,6 +205,20 @@ def test_partition_source_before_removes_newlines():
     assert nl == '\n'
 
 
+def test_partition_source_preserves_whitespace_between_pre_import_and_import():
+    # keep blank lines between a top-level docstring and the first import to
+    # match black's formatting behaviour
+    before, imports, after, nl = partition_source(
+        '"""doc"""\n'
+        '\n'
+        'import os\n',
+    )
+    assert before == '"""doc"""\n\n'
+    assert imports == ['import os\n']
+    assert after == ''
+    assert nl == '\n'
+
+
 def test_partition_source_before_and_code_only():
     before, imports, after, nl = partition_source(
         '# before\n'