Use CSafeLoader when possible (#1150)

Use CSafeLoader when possible --------- Co-authored-by: Olivier Delalleau <507137+odelalleau@users.noreply.github.com>
omry · Jan 16, 2024 · b5badbd · b5badbd
1 parent fd73050
commit b5badbd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/news/1150.feature b/news/1150.feature
@@ -0,0 +1 @@
+The YAML parser will now use `yaml.CSafeLoader` instead of `yaml.SafeLoader` whenever possible to speed up parsing
diff --git a/omegaconf/_utils.py b/omegaconf/_utils.py
@@ -34,6 +34,13 @@
 except ImportError:  # pragma: no cover
     attr = None  # type: ignore # pragma: no cover
 
+try:
+    from yaml import CSafeLoader
+
+    BaseLoader = CSafeLoader
+except ImportError:  # pragma: no cover
+    BaseLoader = yaml.SafeLoader
+
 NoneType: Type[None] = type(None)
 
 BUILTIN_VALUE_TYPES: Tuple[Type[Any], ...] = (
@@ -123,7 +130,7 @@ def yaml_is_bool(b: str) -> bool:
 
 
 def get_yaml_loader() -> Any:
-    class OmegaConfLoader(yaml.SafeLoader):  # type: ignore
+    class OmegaConfLoader(BaseLoader):  # type: ignore
         def construct_mapping(self, node: yaml.Node, deep: bool = False) -> Any:
             keys = set()
             for key_node, value_node in node.value: