Skip to content

Commit

Permalink
simplify add_imports by handling blank-only files earlier
Browse files Browse the repository at this point in the history
  • Loading branch information
asottile committed Jul 3, 2022
1 parent f962c00 commit d123abe
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 10 deletions.
16 changes: 6 additions & 10 deletions reorder_python_imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,17 +153,9 @@ def _inner() -> Generator[CodePartition, None, None]:


def add_imports(
partitions: Iterable[CodePartition],
partitions: list[CodePartition],
to_add: tuple[str, ...] = (),
) -> list[CodePartition]:
partitions = list(partitions)
if not _partitions_to_src(partitions).strip():
return partitions

# If we don't have a trailing newline, this refactor is wrong
if not partitions[-1].src.endswith('\n'):
partitions[-1] = partitions[-1]._replace(src=partitions[-1].src + '\n')

return partitions + [
CodePartition(CodeType.IMPORT, imp_statement.strip() + '\n')
for imp_statement in to_add
Expand Down Expand Up @@ -400,7 +392,11 @@ def fix_file_contents(
) -> str:
# internally use `'\n` as the newline and normalize at the very end
nl = _most_common_line_ending(contents)
contents = contents.replace('\r\n', '\n').replace('\r', '\n')
contents = contents.replace('\r\n', '\n').replace('\r', '\n').rstrip()
if contents:
contents += '\n'
else:
return ''

partitioned = partition_source(contents)
partitioned = combine_trailing_code_chunks(partitioned)
Expand Down
9 changes: 9 additions & 0 deletions tests/reorder_python_imports_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -346,6 +346,15 @@ def test_add_import_trivial():
) == ''


def test_add_imports_empty_file():
assert fix_file_contents(
'\n\n',
to_add=('from __future__ import absolute_import',),
to_remove=set(),
to_replace=Replacements.make([]),
) == ''


def test_add_import_import_already_there():
assert fix_file_contents(
'from __future__ import absolute_import\n',
Expand Down

0 comments on commit d123abe

Please sign in to comment.