Skip to content

Commit

Permalink
feat: allow formatting files with multiple documents
Browse files Browse the repository at this point in the history
  • Loading branch information
muripic committed Aug 21, 2021
1 parent ecc1f4f commit 1f2e301
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 3 deletions.
7 changes: 4 additions & 3 deletions src/yamlfix/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,12 +96,13 @@ def _ruamel_yaml_fixer(source_code: str) -> str:
# Start the document with ---
# ignore: variable has type None, what can we do, it doesn't have type hints...
yaml.explicit_start = True # type: ignore
source_dict = yaml.load(source_code)
source_dicts = yaml.load_all(source_code)

# Return the output to a string
string_stream = StringIO()
yaml.dump(source_dict, string_stream)
source_code = string_stream.getvalue()
for source_dict in source_dicts:
yaml.dump(source_dict, string_stream)
source_code = string_stream.getvalue()
string_stream.close()

return source_code.strip()
Expand Down
17 changes: 17 additions & 0 deletions tests/unit/test_services.py
Original file line number Diff line number Diff line change
Expand Up @@ -339,3 +339,20 @@ def test_fix_code_doesnt_change_double_exclamation_marks() -> None:
result = fix_code(source)

assert result == source


def test_fix_code_parses_files_with_multiple_documents() -> None:
"""Files that contain multiple documents should be parsed as a collection of
separate documents and then dumped together again.
"""
source = dedent(
"""\
---
project: yamlfix
---
project: yamlfix"""
)

result = fix_code(source)

assert result == source

0 comments on commit 1f2e301

Please sign in to comment.