[Mlir-commits] [mlir] [MLIR][Python] Register `OpAttributeMap` as `Mapping` for `match` compatibility (PR #174292)

[llvmlistbot at llvm.org]

https://github.com/MaPePeR updated https://github.com/llvm/llvm-project/pull/174292

>From 86ed86690a0a32fd06373859882a1ad9a6beec2c Mon Sep 17 00:00:00 2001
From: MaPePeR <MaPePeR at users.noreply.github.com>
Date: Wed, 31 Dec 2025 19:03:34 +0000
Subject: [PATCH] [MLIR][Python] Register OpAttributeMap as Mapping for match
 compatibility

---
 mlir/include/mlir/Bindings/Python/IRCore.h |  3 ++
 mlir/lib/Bindings/Python/IRCore.cpp        | 14 +++++
 mlir/python/mlir/_mlir_libs/__init__.py    |  3 +-
 mlir/test/python/ir/operation.py           | 63 ++++++++++++++++++++++
 4 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 8a8d2a1b2270f..b2eb6b9edb677 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1782,6 +1782,9 @@ class MLIR_PYTHON_API_EXPORTED PyOpAttributeMap {
 
   PyNamedAttribute dunderGetItemIndexed(intptr_t index);
 
+  nanobind::object getWithDefaultNamed(const std::string &key,
+                                       nanobind::object defaultValue);
+
   void dunderSetItem(const std::string &name, const PyAttribute &attr);
 
   void dunderDelItem(const std::string &name);
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 8396569757183..eb5ebf15ea728 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2388,6 +2388,16 @@ PyOpAttributeMap::dunderGetItemNamed(const std::string &name) {
   return PyAttribute(operation->getContext(), attr).maybeDownCast();
 }
 
+nb::object PyOpAttributeMap::getWithDefaultNamed(const std::string &key,
+                                                 nb::object defaultValue) {
+  MlirAttribute attr =
+      mlirOperationGetAttributeByName(operation->get(), toMlirStringRef(key));
+  if (mlirAttributeIsNull(attr)) {
+    return defaultValue;
+  }
+  return PyAttribute(operation->getContext(), attr).maybeDownCast();
+}
+
 PyNamedAttribute PyOpAttributeMap::dunderGetItemIndexed(intptr_t index) {
   if (index < 0) {
     index += dunderLen();
@@ -2450,6 +2460,10 @@ void PyOpAttributeMap::bind(nb::module_ &m) {
            "Sets an attribute with the given name.")
       .def("__delitem__", &PyOpAttributeMap::dunderDelItem, "name"_a,
            "Deletes an attribute with the given name.")
+      .def("get", &PyOpAttributeMap::getWithDefaultNamed, nb::arg("key"),
+           nb::arg("default") = nb::none(),
+           "Gets an attribute by name or the default value, if it does not "
+           "exist.")
       .def(
           "__iter__",
           [](PyOpAttributeMap &self) {
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index ce7e6bf93012a..c0e8775149d41 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -2,7 +2,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from typing import Any, Sequence
+from typing import Any, Mapping, Sequence
 
 import os
 
@@ -245,6 +245,7 @@ def __str__(self):
     Sequence.register(ir.OpResultList)
     Sequence.register(ir.OpSuccessors)
     Sequence.register(ir.RegionSequence)
+    Mapping.register(ir.OpAttributeMap)
 
 
 _site_initialize()
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index b9242a7cc2bd9..89f78ab1932a0 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -652,6 +652,34 @@ def testOperationAttributes():
     # CHECK: Dict mapping {'dependent': 'text', 'other.attribute': 3.0, 'some.attribute': 1}
     print("Dict mapping", d)
 
+    # Structural pattern matching test using Mapping
+
+    # CHECK: Matched Mapping Attribute 'some.attribute': 1
+    # CHECK: Matched Mapping Attribute 'other.attribute': 3.0
+    # CHECK: Matched Mapping Attribute 'dependent': text
+    match op:
+        case OpView(attributes={"does_not_exist": a0}):
+            assert False
+        case OpView(
+            attributes={
+                "some.attribute": IntegerAttr(value=some_attr_val) as some_attr,
+                "other.attribute": FloatAttr() as other_attr,
+                "dependent": StringAttr() as dep_attr,
+                **other_attributes,
+            }
+        ):
+            print(f"Matched Mapping Attribute 'some.attribute': {some_attr_val}")
+            print(f"Matched Mapping Attribute 'other.attribute': {other_attr.value}")
+            print(f"Matched Mapping Attribute 'dependent': {dep_attr.value}")
+            assert type(other_attributes) == dict
+            assert len(other_attributes) == 0
+            assert some_attr == op.attributes.get("some.attribute")
+            assert other_attr == op.attributes.get("other.attribute", None)
+            assert dep_attr == op.attributes.get("dependent", "Default value")
+        case _:
+            print("Did not match!")
+            assert False
+
     # Check that exceptions are raised as expected.
     try:
         op.attributes["does_not_exist"]
@@ -667,6 +695,41 @@ def testOperationAttributes():
     else:
         assert False, "expected IndexError on accessing an out-of-bounds attribute"
 
+    # Check that exceptions are raised when `get` is used with non-str arg.
+    try:
+        op.attributes.get(0)
+    except TypeError:
+        pass
+    else:
+        assert False, "expected TypeError using int as key for get()"
+
+    try:
+        op.attributes.get(0, None)
+    except TypeError:
+        pass
+    else:
+        assert False, "expected TypeError using int as key for get()"
+
+    try:
+        op.attributes.get([], None)
+    except TypeError:
+        pass
+    else:
+        assert False, "expected TypeError using list as key for get()"
+
+    try:
+        match op:
+            case OpView(attributes={0: a}):
+                assert False
+    except TypeError:
+        pass
+    else:
+        assert False, "expected TypeError matching OpAttributeMap with int-key "
+
+    # get() does not throw for non existent attributes.
+    assert op.attributes.get("does_not_exist") is None
+    assert op.attributes.get("does_not_exist", "default_value") == "default_value"
+
 
 # CHECK-LABEL: TEST: testOperationPrint
 @run