diff --git a/reorder_python_imports.py b/reorder_python_imports.py index c6d1f57..6c9125c 100644 --- a/reorder_python_imports.py +++ b/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: diff --git a/tests/reorder_python_imports_test.py b/tests/reorder_python_imports_test.py index c61c39a..e1b9b8f 100644 --- a/tests/reorder_python_imports_test.py +++ b/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'