update yaml safe loading

datajoint · Dec 6, 2023 · 7c8ccdd · 7c8ccdd
1 parent 556556b
commit 7c8ccdd
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/element_deeplabcut/model.py b/element_deeplabcut/model.py
@@ -7,7 +7,7 @@
 import os
 import cv2
 import csv
-import yaml
+from ruamel.yaml import YAML
 import inspect
 import importlib
 import numpy as np
@@ -239,8 +239,9 @@ def extract_new_body_parts(cls, dlc_config: dict, verbose: bool = True):
                 ".yaml",
             ), f"dlc_config is neither dict nor filepath\n Check: {dlc_config_fp}"
             if dlc_config_fp.suffix in (".yml", ".yaml"):
+                yaml = YAML(typ="safe", pure=True)
                 with open(dlc_config_fp, "rb") as f:
-                    dlc_config = yaml.safe_load(f)
+                    dlc_config = yaml.load(f)
         # -- Check and insert new BodyPart --
         assert "bodyparts" in dlc_config, f"Found no bodyparts section in {dlc_config}"
         tracked_body_parts = cls.fetch("body_part")
@@ -381,8 +382,9 @@ def insert_new_model(
             "dlc_config is not a filepath" + f"\n Check: {dlc_config_fp}"
         )
         if dlc_config_fp.suffix in (".yml", ".yaml"):
+            yaml = YAML(typ="safe", pure=True)
             with open(dlc_config_fp, "rb") as f:
-                dlc_config = yaml.safe_load(f)
+                dlc_config = yaml.load(f)
         if isinstance(params, dict):
             dlc_config.update(params)