fix(plan reading): use yaml safe_load

Signed-off-by: Pant, Akshay <[email protected]>
securefederatedai · Mar 7, 2025 · b2f592e · b2f592e
1 parent 1c3aaba
commit b2f592e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/openfl/experimental/workflow/interface/cli/aggregator.py b/openfl/experimental/workflow/interface/cli/aggregator.py
@@ -12,7 +12,6 @@
 import yaml
 from click import Path as ClickPath
 from click import confirm, echo, group, option, pass_context, style
-from yaml.loader import SafeLoader
 
 from openfl.cryptography.ca import sign_certificate
 from openfl.cryptography.io import get_csr_hash, read_crt, read_csr, read_key, write_crt, write_key
@@ -80,7 +79,7 @@ def start_(plan, authorized_cols, secure):
         )
     else:
         with open("plan/data.yaml", "r") as f:
-            data = yaml.load(f, Loader=SafeLoader)
+            data = yaml.safe_load(f)
             if data.get("aggregator", None) is None:
                 logger.warning(
                     "Aggregator private attributes are set to None as no aggregator"

diff --git a/openfl/experimental/workflow/interface/cli/collaborator.py b/openfl/experimental/workflow/interface/cli/collaborator.py
@@ -18,7 +18,6 @@
 from click import Path as ClickPath
 from click import confirm, echo, group, option, pass_context, style
 from yaml import FullLoader, dump, load
-from yaml.loader import SafeLoader
 
 from openfl.cryptography.ca import sign_certificate
 from openfl.cryptography.io import get_csr_hash, read_crt, read_csr, read_key, write_crt, write_key
@@ -83,7 +82,7 @@ def start_(plan, collaborator_name, secure, data_config="plan/data.yaml"):
         )
     else:
         with open(data_config, "r") as f:
-            data = yaml.load(f, Loader=SafeLoader)
+            data = yaml.safe_load(f)
             if data.get(collaborator_name, None) is None:
                 logger.warning(
                     f"Collaborator private attributes are set to None as no attributes"