Merge pull request #112 from muripic/feat/format-multiple-doc-files

feat: Allow formatting files with multiple documents

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
       - id: check-merge-conflict
       - id: end-of-file-fixer
   - repo: https://github.com/ambv/black
-    rev: master
+    rev: main
     hooks:
       - id: black
         language_version: python3.7

src/yamlfix/services.py

@@ -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()

tests/unit/test_services.py

@@ -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