From 1f2e3010d6c82169d3ab4fd5fec312d89545b5fd Mon Sep 17 00:00:00 2001
From: muripic <muri.pic@gmail.com>
Date: Sat, 21 Aug 2021 11:09:53 -0300
Subject: [PATCH] feat: allow formatting files with multiple documents

---
 src/yamlfix/services.py     |  7 ++++---
 tests/unit/test_services.py | 17 +++++++++++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/yamlfix/services.py b/src/yamlfix/services.py
index 32fcc3a..59d0992 100644
--- a/src/yamlfix/services.py
+++ b/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()
diff --git a/tests/unit/test_services.py b/tests/unit/test_services.py
index ed0e6d1..a5f11de 100644
--- a/tests/unit/test_services.py
+++ b/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