yaml supports include dirs

1 files changed, 27 insertions, 3 deletions

diff --git a/kvmd/yamlconf/loader.py b/kvmd/yamlconf/loader.py
index bc38d7c1..381a8c4f 100644
--- a/kvmd/yamlconf/loader.py
+++ b/kvmd/yamlconf/loader.py
@@ -22,6 +22,8 @@
 
 import os
 
+from typing import List
+from typing import Dict
 from typing import IO
 from typing import Any
 
@@ -49,9 +51,31 @@ class _YamlLoader(yaml.SafeLoader):
         super().__init__(yaml_file)
         self.__root = os.path.dirname(yaml_file.name)
 
-    def include(self, node: yaml.nodes.ScalarNode) -> Any:
-        path = os.path.join(self.__root, str(self.construct_scalar(node)))
-        return load_yaml_file(path)
+    def include(self, node: yaml.nodes.Node) -> Any:
+        incs: List[str]
+        if isinstance(node, yaml.nodes.SequenceNode):
+            incs = [
+                str(child)
+                for child in self.construct_sequence(node)
+                if isinstance(child, (int, float, str))
+            ]
+        else:  # Trying scalar for the fallback
+            incs = [str(self.construct_scalar(node))]  # type: ignore
+        return self.__inner_include(list(filter(None, incs)))
+
+    def __inner_include(self, incs: List[str]) -> Any:
+        tree: Dict = {}
+        for inc in filter(None, incs):
+            assert inc, inc
+            inc_path = os.path.join(self.__root, inc)
+            if os.path.isdir(inc_path):
+                for child in sorted(os.listdir(inc_path)):
+                    child_path = os.path.join(inc_path, child)
+                    if os.path.isfile(child_path) or os.path.islink(child_path):
+                        tools.merge(tree, (load_yaml_file(child_path) or {}))
+            else:  # Try file
+                tools.merge(tree, (load_yaml_file(inc_path) or {}))
+        return tree
 
 
 _YamlLoader.add_constructor("!include", _YamlLoader.include)