From 2a8764577f693d11ca405c7aa2b5a4296a5323c0 Mon Sep 17 00:00:00 2001
From: ldericher <40151420+ldericher@users.noreply.github.com>
Date: Thu, 28 Oct 2021 13:54:02 +0200
Subject: [PATCH] misc.YAML.dump* methods

---
 kiwi_scp/config.py   | 13 ++++---------
 kiwi_scp/misc.py     | 30 ++++++++++++++++++++++--------
 tests/test_config.py |  9 +--------
 3 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/kiwi_scp/config.py b/kiwi_scp/config.py
index e4bf1b9..48a8217 100644
--- a/kiwi_scp/config.py
+++ b/kiwi_scp/config.py
@@ -7,7 +7,7 @@ from typing import Optional, Dict, List, Any, TextIO
 from pydantic import BaseModel, constr, root_validator, validator
 
 from ._constants import RE_SEMVER, RE_VARNAME, KIWI_CONF_NAME
-from .misc import YAML, _format_kiwi_yml
+from .misc import YAML
 
 
 class StorageConfig(BaseModel):
@@ -187,21 +187,16 @@ class KiwiConfig(BaseModel):
 
         return result
 
-    def dump_kiwi_yml(self, stream: TextIO) -> None:
+    def dump_kiwi_yml(self, stream: TextIO = None) -> Optional[str]:
         """dump a kiwi.yml file"""
 
-        YAML().dump(self.kiwi_dict, stream=stream, transform=_format_kiwi_yml)
+        return YAML().dump_kiwi_yml(self.kiwi_dict, stream=stream)
 
     @property
     def kiwi_yml(self) -> str:
         """get a kiwi.yml dump as a string"""
 
-        sio = io.StringIO()
-        self.dump_kiwi_yml(sio)
-        result: str = sio.getvalue()
-        sio.close()
-
-        return result
+        return self.dump_kiwi_yml()
 
     @validator("shells", pre=True)
     @classmethod
diff --git a/kiwi_scp/misc.py b/kiwi_scp/misc.py
index 72d5d8c..c4895f0 100644
--- a/kiwi_scp/misc.py
+++ b/kiwi_scp/misc.py
@@ -1,9 +1,10 @@
 import re
-from typing import Any, Type, List, Callable
+from typing import Any, Type, List, Callable, Optional
 
 import attr
 import click
 import ruamel.yaml
+import ruamel.yaml.compat
 from click.decorators import FC
 
 from ._constants import HEADER_KIWI_CONF_NAME
@@ -60,16 +61,29 @@ class YAML(ruamel.yaml.YAML):
         super().__init__(*args, **kwargs)
         self.indent(offset=2)
 
+    def dump(self, data, stream=None, **kwargs) -> Optional[str]:
+        into_str: bool = False
+        if stream is None:
+            into_str = True
+            stream = ruamel.yaml.compat.StringIO()
 
-def _format_kiwi_yml(yml_string: str):
-    # insert newline before every main key
-    yml_string = re.sub(r'^(\S)', r'\n\1', yml_string, flags=re.MULTILINE)
+        super().dump(data, stream=stream, **kwargs)
+        if into_str:
+            return stream.getvalue()
 
-    # load header comment from file
-    with open(HEADER_KIWI_CONF_NAME, 'r') as stream:
-        yml_string = stream.read() + yml_string
+    @staticmethod
+    def _format_kiwi_yml(yml_string: str):
+        # insert newline before every main key
+        yml_string = re.sub(r'^(\S)', r'\n\1', yml_string, flags=re.MULTILINE)
 
-    return yml_string
+        # load header comment from file
+        with open(HEADER_KIWI_CONF_NAME, 'r') as stream:
+            yml_string = stream.read() + yml_string
+
+        return yml_string
+
+    def dump_kiwi_yml(self, data, **kwargs) -> Optional[str]:
+        return self.dump(data, transform=YAML._format_kiwi_yml, **kwargs)
 
 
 def _surround(string, bang):
diff --git a/tests/test_config.py b/tests/test_config.py
index 1d09e05..61b78a1 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,4 +1,3 @@
-import io
 from ipaddress import IPv4Network
 from pathlib import Path
 
@@ -45,13 +44,7 @@ class TestDefault:
         }
         assert c.kiwi_dict == kiwi_dict
 
-        sio = io.StringIO()
-        from kiwi_scp.misc import _format_kiwi_yml
-        YAML(typ="safe").dump(kiwi_dict, stream=sio, transform=_format_kiwi_yml)
-        yml_string = sio.getvalue()
-        sio.close()
-
-        assert c.kiwi_yml == yml_string
+        assert c.kiwi_yml == YAML().dump_kiwi_yml(kiwi_dict)
 
 
 class TestVersion: