Keep whitespace before imports

This commits comes from

  * asottile#370
  * psf/black#4175 (comment)

because

  * asottile#367
  * asottile#366
  * psf/black#4175

Suggested-by: Achim Herwig <achim.herwig@wodca.de>

reorder_python_imports.py

@@ -94,10 +94,10 @@ def partition_source(src: str) -> tuple[str, list[str], str, str]:
             pre_import = False
             chunks.append((CodeType.IMPORT, s))
         elif token_type is Tok.NEWLINE:
-            if s.isspace():
-                tp = CodeType.NON_CODE
-            elif pre_import:
+            if pre_import:
                 tp = CodeType.PRE_IMPORT_CODE
+            elif s.isspace():
+                tp = CodeType.NON_CODE
             else:
                 tp = CodeType.CODE
 

tests/reorder_python_imports_test.py

@@ -47,7 +47,7 @@ def test_tokenize_can_match_strings(s):
 @pytest.mark.parametrize(
     's',
     (
-        pytest.param('', id='trivial'),
+        pytest.param('\n', id='trivial'),
         pytest.param('#!/usr/bin/env python\n', id='shebang'),
         pytest.param('# -*- coding: UTF-8 -*-\n', id='source encoding'),
         pytest.param('  # coding: UTF-8\n', id='source encoding indented'),
@@ -190,17 +190,21 @@ def test_partition_source_imports_only(s, expected):
     assert nl == '\n'
 
 
-def test_partition_source_before_removes_newlines():
+def test_partition_source_before_leaves_newlines():
     before, imports, after, nl = partition_source(
         '# comment here\n'
         '\n'
-        '# another comment here\n',
+        '# another comment here\n'
+        '\n'
+        'import os\n'
     )
     assert before == (
         '# comment here\n'
+        '\n'
         '# another comment here\n'
+        '\n'
     )
-    assert imports == []
+    assert imports == ['import os\n']
     assert after == ''
     assert nl == '\n'