Fix unsafe YAML config loading

Author: isVoid

numbast/src/numbast/experimental/mlir/tools/static_binding_generator.py

@@ -45,7 +45,7 @@
 VERBOSE = True
 
 # Register custom YAML constructor for !join tag
-yaml.add_constructor("!numbast_join", string_constructor)
+yaml.SafeLoader.add_constructor("!numbast_join", string_constructor)
 
 
 def _config_dict_uses_mlir_backend(config_dict: dict) -> bool:
@@ -59,7 +59,7 @@ def _config_dict_uses_mlir_backend(config_dict: dict) -> bool:
 
 def _cfg_path_uses_mlir_backend(cfg_path: str) -> bool:
     with open(cfg_path) as f:
-        config_dict = yaml.load(f, yaml.Loader)
+        config_dict = yaml.safe_load(f)
     return _config_dict_uses_mlir_backend(config_dict)
 
 
@@ -262,7 +262,7 @@ def from_yaml_path(cls, cfg_path: str) -> "Config":
             A new Config instance.
         """
         with open(cfg_path) as f:
-            config_dict = yaml.load(f, yaml.Loader)
+            config_dict = yaml.safe_load(f)
         return cls(config_dict)
 
     @classmethod

numbast/src/numbast/experimental/mlir/tools/tests/test_config_loading.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import yaml
+
+from numbast.experimental.mlir.tools import static_binding_generator as sbg
+
+
+@pytest.mark.parametrize(
+    "load_config",
+    [sbg._cfg_path_uses_mlir_backend, sbg.Config.from_yaml_path],
+)
+def test_config_load_rejects_python_object_tags(tmp_path, load_config):
+    marker_path = tmp_path / "unsafe-loader-marker"
+    cfg_path = tmp_path / "config.yaml"
+    cfg_path.write_text(
+        f"!!python/object/apply:builtins.open\n- {marker_path}\n- w\n",
+        encoding="utf-8",
+    )
+
+    with pytest.raises(yaml.constructor.ConstructorError):
+        load_config(cfg_path)
+
+    assert not marker_path.exists()

numbast/src/numbast/tools/static_binding_generator.py

@@ -52,7 +52,7 @@
 )
 
 # Register custom YAML constructor for !join tag
-yaml.add_constructor("!numbast_join", string_constructor)
+yaml.SafeLoader.add_constructor("!numbast_join", string_constructor)
 
 
 def _config_dict_uses_mlir_backend(config_dict: dict) -> bool:
@@ -100,7 +100,7 @@ def _validate_mlir_backend_only_config(config_dict: dict):
 
 def _cfg_path_uses_mlir_backend(cfg_path: str) -> bool:
     with open(cfg_path) as f:
-        config_dict = yaml.load(f, yaml.Loader)
+        config_dict = yaml.safe_load(f)
     return _config_dict_uses_mlir_backend(config_dict)
 
 
@@ -227,7 +227,7 @@ def from_yaml_path(cls, cfg_path: str) -> "Config":
             A new Config instance.
         """
         with open(cfg_path) as f:
-            config_dict = yaml.load(f, yaml.Loader)
+            config_dict = yaml.safe_load(f)
         return cls(config_dict)
 
     @classmethod

numbast/src/numbast/tools/tests/test_mlir_backend_routing.py

@@ -26,6 +26,34 @@ def test_cfg_path_uses_mlir_backend(tmp_path):
     assert sbg._cfg_path_uses_mlir_backend(cfg_path)
 
 
+def test_cfg_path_supports_numbast_join_tag(tmp_path):
+    cfg_path = tmp_path / "config.yaml"
+    cfg_path.write_text(
+        'MLIR Backend: !numbast_join ["tr", "ue"]\n',
+        encoding="utf-8",
+    )
+
+    assert sbg._cfg_path_uses_mlir_backend(cfg_path)
+
+
+@pytest.mark.parametrize(
+    "load_config",
+    [sbg._cfg_path_uses_mlir_backend, sbg.Config.from_yaml_path],
+)
+def test_config_load_rejects_python_object_tags(tmp_path, load_config):
+    marker_path = tmp_path / "unsafe-loader-marker"
+    cfg_path = tmp_path / "config.yaml"
+    cfg_path.write_text(
+        f"!!python/object/apply:builtins.open\n- {marker_path}\n- w\n",
+        encoding="utf-8",
+    )
+
+    with pytest.raises(yaml.constructor.ConstructorError):
+        load_config(cfg_path)
+
+    assert not marker_path.exists()
+
+
 def test_static_generator_dispatches_mlir_backend(monkeypatch, tmp_path):
     class DummyConfig:
         mlir_backend = True