From 06fd9aa6e1f42ca31dc353417e1cb225253c3d31 Mon Sep 17 00:00:00 2001 From: Anthony Sottile Date: Sun, 3 Jul 2022 09:01:26 -0700 Subject: [PATCH] combine combine_trailing_code_chunks and partition_source --- reorder_python_imports.py | 30 +++++++++++------------------- 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/reorder_python_imports.py b/reorder_python_imports.py index 86dc4c9..6f96a8b 100644 --- a/reorder_python_imports.py +++ b/reorder_python_imports.py @@ -29,6 +29,8 @@ CodeType = enum.Enum('CodeType', 'PRE_IMPORT_CODE IMPORT NON_CODE CODE') +NON_COMBINABLE = (CodeType.IMPORT, CodeType.PRE_IMPORT_CODE) + class CodePartition(NamedTuple): code_type: CodeType @@ -45,9 +47,6 @@ def _partitions_to_src(partitions: Iterable[CodePartition]) -> str: def partition_source(src: str) -> list[CodePartition]: - """Partitions source into a list of `CodePartition`s for import - refactoring. - """ lines = src.splitlines(True) chunks = [] @@ -115,22 +114,16 @@ def partition_source(src: str) -> list[CodePartition]: chunks.append(CodePartition(CodeType.CODE, srctext)) break - return [chunk for chunk in chunks if chunk.src] - - -def combine_trailing_code_chunks( - partitions: Iterable[CodePartition], -) -> list[CodePartition]: - chunks = list(partitions) - - NON_COMBINABLE = (CodeType.IMPORT, CodeType.PRE_IMPORT_CODE) - if chunks and chunks[-1].code_type not in NON_COMBINABLE: - src = chunks.pop().src - while chunks and chunks[-1].code_type not in NON_COMBINABLE: - src = chunks.pop().src + src + # combine any trailing code chunks + ret = [chunk for chunk in chunks if chunk.src] + code = [] + while ret and ret[-1].code_type not in NON_COMBINABLE: + code.append(ret.pop()) + if code: + code_src = _partitions_to_src(reversed(code)) + ret.append(CodePartition(CodeType.CODE, code_src)) - chunks.append(CodePartition(CodeType.CODE, src)) - return chunks + return ret def separate_comma_imports( @@ -387,7 +380,6 @@ def fix_file_contents( return '' partitioned = partition_source(contents) - partitioned = combine_trailing_code_chunks(partitioned) partitioned = add_imports(partitioned, to_add=to_add) partitioned = separate_comma_imports(partitioned) partitioned = replace_imports(partitioned, to_replace=to_replace)