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