Fix some uncaught yaml loader codes. Replace them with safe_load (#383)

NVIDIA · Apr 1, 2022 · 71c061f · 71c061f
1 parent 4de9782
commit 71c061f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/nvflare/ha/overseer/utils.py b/nvflare/ha/overseer/utils.py
@@ -16,7 +16,7 @@
 import uuid
 from datetime import datetime, timedelta
 
-import yaml
+from nvflare.lighter.utils import load_yaml
 
 OVERSEER_STORE = os.environ.get("OVERSEER_STORE")
 
@@ -34,7 +34,7 @@
 def load_privilege():
     privilege_file = os.environ.get("AUTHZ_FILE", "privilege.yml")
     try:
-        privilege = yaml.load(open(privilege_file, "tr"), Loader=yaml.Loader)
+        privilege = load_yaml(privilege_file)
     except:
         privilege = dict()
     return privilege

diff --git a/nvflare/lighter/study.py b/nvflare/lighter/study.py
@@ -19,7 +19,7 @@
 import os
 from datetime import datetime
 
-import yaml
+from nvflare.lighter.utils import load_yaml
 
 
 def get_input(prompt, item_list, multiple=False):
@@ -75,7 +75,7 @@ def main():
         print(f"{project_full_path} not found.  Running study requires that file.")
         exit(0)
 
-    project = yaml.load(open(project_full_path, "r"), Loader=yaml.Loader)
+    project = load_yaml(project_full_path)
     api_version = project.get("api_version")
     if api_version not in [3]:
         raise ValueError(f"API version expected 3 but found {api_version}")