Add yaml include support

nvflare/lighter/utils.py

@@ -181,12 +181,40 @@ def sign_all(content_folder, signing_pri_key):
 
 
 def load_yaml(file):
+
+    root = os.path.split(file)[0]
+    yaml_data = None
     if isinstance(file, str):
-        return yaml.safe_load(open(file, "r"))
+        yaml_data = yaml.safe_load(open(file, "r"))
     elif isinstance(file, bytes):
-        return yaml.safe_load(file)
-    else:
-        return None
+        yaml_data = yaml.safe_load(file)
+
+    yaml_data = load_yaml_include(root, yaml_data)
+
+    return yaml_data
+
+
+def load_yaml_include(root, yaml_data):
+    new_data = {}
+    for k, v in yaml_data.items():
+        if k == "include":
+            if isinstance(v, str):
+                includes = [v]
+            elif isinstance(v, list):
+                includes = v
+            for item in includes:
+                new_data.update(load_yaml(os.path.join(root, item)))
+        elif isinstance(v, list):
+            new_list = []
+            for item in v:
+                if isinstance(item, dict):
+                    item = load_yaml_include(root, item)
+                new_list.append(item)
+            new_data[k] = new_list
+        else:
+            new_data[k] = v
+
+    return new_data
 
 
 def sh_replace(src, mapping_dict):

tests/unit_test/lighter/0.yml

@@ -0,0 +1,12 @@
+api_version: 3
+name: example_project
+
+include: 1.yml
+
+participants:
+  - name: server
+    port: 123
+    include: [1.yml]
+  - name: client
+    port: 234
+    include: 2.yml

tests/unit_test/lighter/1.yml

@@ -0,0 +1 @@
+server_name: server

tests/unit_test/lighter/2.yml

@@ -0,0 +1 @@
+client_name: client-1

tests/unit_test/lighter/utils_test.py

@@ -25,7 +25,7 @@
 from cryptography.x509.oid import NameOID
 
 from nvflare.lighter.impl.cert import serialize_cert
-from nvflare.lighter.utils import sign_folders, verify_folder_signature
+from nvflare.lighter.utils import load_yaml, sign_folders, verify_folder_signature
 
 folders = ["folder1", "folder2"]
 files = ["file1", "file2"]
@@ -144,3 +144,20 @@ def test_verify_updated_folder(self):
         os.unlink("client.crt")
         os.unlink("root.crt")
         shutil.rmtree(folder)
+
+    def _get_participant(self, name, participants):
+        for p in participants:
+            if p.get("name") == name:
+                return p
+
+    def test_load_yaml(self):
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        data = load_yaml(os.path.join(dir_path, "0.yml"))
+
+        assert data.get("server_name") == "server"
+
+        participant = self._get_participant("server", data.get("participants"))
+        assert participant.get("server_name") == "server"
+
+        participant = self._get_participant("client", data.get("participants"))
+        assert participant.get("client_name") == "client-1"