diff options
Diffstat (limited to 'kvmd/yamlconf/loader.py')
| -rw-r--r-- | kvmd/yamlconf/loader.py | 30 |
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) |