From 730324f1452032b4b739339dd19e2b3e587e44da Mon Sep 17 00:00:00 2001
From: Felix Fontein <felix@fontein.de>
Date: Fri, 30 Aug 2024 14:21:43 +0200
Subject: [PATCH] Add extended YAML functionality from antsibull-changelog (#2)

---
 src/antsibull_fileutils/yaml.py | 31 ++++++++++++++++++++-----
 tests/units/test_yaml.py        | 40 ++++++++++++++++++++++++++++-----
 2 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/src/antsibull_fileutils/yaml.py b/src/antsibull_fileutils/yaml.py
index 6017900..5de9bae 100644
--- a/src/antsibull_fileutils/yaml.py
+++ b/src/antsibull_fileutils/yaml.py
@@ -20,7 +20,7 @@
     # use C version if possible for speedup
     from yaml import CSafeDumper as _SafeDumper
     from yaml import CSafeLoader as _SafeLoader
-except ImportError:
+except ImportError:  # pragma: no cover
     from yaml import SafeDumper as _SafeDumper
     from yaml import SafeLoader as _SafeLoader
 
@@ -28,6 +28,15 @@
     from _typeshed import StrOrBytesPath, SupportsWrite
 
 
+class _IndentedDumper(yaml.SafeDumper):
+    """
+    Extend YAML dumper to increase indent of list items.
+    """
+
+    def increase_indent(self, flow=False, indentless=False):
+        return super().increase_indent(flow, False)
+
+
 def load_yaml_bytes(data: bytes) -> t.Any:
     """
     Load and parse YAML from given bytes.
@@ -43,20 +52,30 @@ def load_yaml_file(path: StrOrBytesPath) -> t.Any:
         return yaml.load(stream, Loader=_SafeLoader)
 
 
-def store_yaml_file(path: StrOrBytesPath, content: t.Any) -> None:
+def store_yaml_file(
+    path: StrOrBytesPath, content: t.Any, *, nice: bool = False, sort_keys: bool = True
+) -> None:
     """
     Store ``content`` as YAML file under ``path``.
     """
     with open(path, "wb") as stream:
-        store_yaml_stream(stream, content)
+        store_yaml_stream(stream, content, nice=nice, sort_keys=sort_keys)
 
 
-def store_yaml_stream(stream: SupportsWrite, content: t.Any) -> None:
+def store_yaml_stream(
+    stream: SupportsWrite, content: t.Any, *, nice: bool = False, sort_keys: bool = True
+) -> None:
     """
     Dump ``content`` as YAML to an IO ``stream``.
     """
-    dumper = _SafeDumper
+    dumper = _IndentedDumper if nice else _SafeDumper
     dumper.ignore_aliases = lambda *args: True
     yaml.dump(
-        content, stream, default_flow_style=False, encoding="utf-8", Dumper=dumper
+        content,
+        stream,
+        default_flow_style=False,
+        encoding="utf-8",
+        Dumper=dumper,
+        explicit_start=nice,
+        sort_keys=sort_keys,
     )
diff --git a/tests/units/test_yaml.py b/tests/units/test_yaml.py
index fe599b7..3582b5e 100644
--- a/tests/units/test_yaml.py
+++ b/tests/units/test_yaml.py
@@ -43,30 +43,58 @@ def test_load_yaml(content, expected, tmp_path):
 
 STORE_YAML_DATA = [
     (
+        {"foo": "bar", "baz": ["bam", "foobar", "aaa"]},
+        {},
+        """baz:
+- bam
+- foobar
+- aaa
+foo: bar
+""",
+    ),
+    (
+        {"foo": "bar", "baz": ["bam", "foobar", "aaa"]},
         {
-            "foo": "bar",
+            "nice": True,
+        },
+        """---
+baz:
+  - bam
+  - foobar
+  - aaa
+foo: bar
+""",
+    ),
+    (
+        {"foo": "bar", "baz": ["bam", "foobar", "aaa"]},
+        {
+            "sort_keys": False,
         },
         """foo: bar
+baz:
+- bam
+- foobar
+- aaa
 """,
     ),
 ]
 
 
 @pytest.mark.parametrize(
-    "content, expected",
+    "content, kwargs, expected",
     STORE_YAML_DATA,
 )
-def test_store_yaml(content, expected, tmp_path):
+def test_store_yaml(content, kwargs, expected, tmp_path):
     file = tmp_path / "test.yaml"
 
-    store_yaml_file(file, content)
+    store_yaml_file(file, content, **kwargs)
     with file.open("r", encoding="utf-8") as f:
         assert f.read() == expected
 
-    store_yaml_file(str(file), content)
+    store_yaml_file(str(file), content, **kwargs)
     with file.open("r", encoding="utf-8") as f:
         assert f.read() == expected
 
     stream = BytesIO()
-    store_yaml_stream(stream, content)
+    store_yaml_stream(stream, content, **kwargs)
     assert stream.getvalue().decode("utf-8") == expected