[MLIR][Python] remove PyYAML as a dep

Author: makslevental

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",
+]