[Mlir-commits] [llvm] [mlir] [MLIR][Python] remove PyYAML as a dep (PR #169145)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/169145

>From 0b42c3d45d7cc0730672e3172cef93b3e74e4330 Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Sat, 22 Nov 2025 00:29:58 -0500
Subject: [PATCH] [MLIR][Python] remove PyYAML as a dep

---
 .ci/all_requirements.txt                      |  6 +-
 .../dialects/linalg/opdsl/lang/yaml_helper.py | 66 ++++++++++++-------
 2 files changed, 46 insertions(+), 26 deletions(-)

diff --git a/.ci/all_requirements.txt b/.ci/all_requirements.txt
index 4918d7519291f..d6457c1e50ea0 100644
--- a/.ci/all_requirements.txt
+++ b/.ci/all_requirements.txt
@@ -341,7 +341,7 @@ ml-dtypes==0.5.1 ; python_version < "3.13" \
 nanobind==2.9.2 \
     --hash=sha256:c37957ffd5eac7eda349cff3622ecd32e5ee1244ecc912c99b5bc8188bafd16e \
     --hash=sha256:e7608472de99d375759814cab3e2c94aba3f9ec80e62cfef8ced495ca5c27d6e
-    # via -r mlir/python/requirements.txt
+    # via -r mlir/python/requirements-dev.txt
 numpy==2.0.2 \
     --hash=sha256:0123ffdaa88fa4ab64835dcbde75dcdf89c453c922f18dced6e27c90d1d0ec5a \
     --hash=sha256:11a76c372d1d37437857280aa142086476136a8c0f373b2e648ab2c8f18fb195 \
@@ -446,7 +446,7 @@ pyasn1-modules==0.4.2 \
 pybind11==2.13.6 \
     --hash=sha256:237c41e29157b962835d356b370ededd57594a26d5894a795960f0047cb5caf5 \
     --hash=sha256:ba6af10348c12b24e92fa086b39cfba0eff619b61ac77c406167d813b096d39a
-    # via -r mlir/python/requirements.txt
+    # via -r mlir/python/requirements-dev.txt
 pycparser==2.23 \
     --hash=sha256:78816d4f24add8f10a06d6f05b4d424ad9e96cfebf68a4ddc99c65c0720d00c2 \
     --hash=sha256:e5c6e8d3fbad53479cab09ac03729e0a9faf2bee3db8208a550daf5af81a5934
@@ -540,7 +540,7 @@ pyyaml==6.0.1 \
     --hash=sha256:fca0e3a251908a499833aa292323f32437106001d436eca0e6e7833256674585 \
     --hash=sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d \
     --hash=sha256:fd66fc5d0da6d9815ba2cebeb4205f95818ff4b79c3ebe268e75d961704af52f
-    # via -r mlir/python/requirements.txt
+    # via -r mlir/python/requirements-dev.txt
 requests==2.32.5 \
     --hash=sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6 \
     --hash=sha256:dbba0bac56e100853db0ea71b82b4dfd5fe2bf6d3754a8893c3af500cec7d7cf
diff --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py
index 1672656b3a1f8..f3ad92d782eff 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py
@@ -4,23 +4,30 @@
 """YAML serialization is routed through here to centralize common logic."""
 
 import sys
+import warnings
+
+
+def multiline_str_representer(dumper, data):
+    if len(data.splitlines()) > 1:
+        return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
+    else:
+        return dumper.represent_scalar("tag:yaml.org,2002:str", data)
+
 
 try:
-    import yaml
+    from yaml import YAMLObject, add_representer
+
+    add_representer(str, multiline_str_representer)
 except ModuleNotFoundError as e:
-    raise ModuleNotFoundError(
-        f"This tool requires PyYAML but it was not installed. "
-        f"Recommend: {sys.executable} -m pip install PyYAML"
-    ) from e
+    warnings.warn(
+        "PyYAML is not installed; you will not be able to generate linalg yaml definitions."
+    )
 
-__all__ = [
-    "yaml_dump",
-    "yaml_dump_all",
-    "YAMLObject",
-]
+    class YAMLObject:
+        pass
 
 
-class YAMLObject(yaml.YAMLObject):
+class YAMLObject(YAMLObject):
     @classmethod
     def to_yaml(cls, dumper, self):
         """Default to a custom dictionary mapping."""
@@ -33,21 +40,34 @@ def as_linalg_yaml(self):
         return yaml_dump(self)
 
 
-def multiline_str_representer(dumper, data):
-    if len(data.splitlines()) > 1:
-        return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
-    else:
-        return dumper.represent_scalar("tag:yaml.org,2002:str", data)
+def yaml_dump(data, sort_keys=False, **kwargs):
+    try:
+        import yaml
 
+        return yaml.dump(data, sort_keys=sort_keys, **kwargs)
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(
+            f"This tool requires PyYAML but it was not installed. "
+            f"Recommend: {sys.executable} -m pip install PyYAML"
+        ) from e
 
-yaml.add_representer(str, multiline_str_representer)
 
+def yaml_dump_all(data, sort_keys=False, explicit_start=True, **kwargs):
+    try:
+        import yaml
 
-def yaml_dump(data, sort_keys=False, **kwargs):
-    return yaml.dump(data, sort_keys=sort_keys, **kwargs)
+        return yaml.dump_all(
+            data, sort_keys=sort_keys, explicit_start=explicit_start, **kwargs
+        )
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(
+            f"This tool requires PyYAML but it was not installed. "
+            f"Recommend: {sys.executable} -m pip install PyYAML"
+        ) from e
 
 
-def yaml_dump_all(data, sort_keys=False, explicit_start=True, **kwargs):
-    return yaml.dump_all(
-        data, sort_keys=sort_keys, explicit_start=explicit_start, **kwargs
-    )
+__all__ = [
+    "yaml_dump",
+    "yaml_dump_all",
+    "YAMLObject",
+]