MaPePeR 524fde8a4d
[MLIR][Python] Register OpAttributeMap as Mapping for match compatibility (#174292)
This is a continuation of the idea from #174091 to add `match` support
for MLIR containers. In this PR the `OpAttributeMap` container is
registered as a `Mapping`, so be mapped as a "dictionary" in `match`
statements.

For this to work the `get(key, default=None)` method had to be
implemented. Those are pretty much copys of `dunderGetItemNamed` and
`dunderGetItemIndexed` with an added argument and `nb::object` as return
type, because they can now return other types than just `PyAttribute`.
Was unsure if I should refactor this to make `dunderGetItem...` use the
new `getWithDefault...` or if a separate method is preferred. Kept it as
a copy for simplicitys sake for now.

Even though the `OpAttributeMap` supports indexing by `int` and `str`,
Python does not allow to register it as a `Sequence` and a `Mapping` at
the same time. If it is registered as a Sequence it only returns the
attribute names as string, not as `NamedAttribute`. It is technically
possible to also use integer keys for the `dict`-like match, but it
doesn't provide any constraints on the number of attributes, etc., so
probably not recommended.

<details><summary>Example</summary>

```python
from mlir.ir import Context, Module, OpAttributeMap
from collections.abc import Sequence

ctx = Context()
ctx.allow_unregistered_dialects = True
module = Module.parse(
    r"""
"some.op"() { some.attribute = 1 : i8,
                other.attribute = 3.0,
                dependent = "text" } : () -> ()
""",
    ctx,
)
op = module.body.operations[0]

def test(attr):
    match attr:
        case [*args]:
            print("matched a Sequence", args)
        case _:
            print("Didn't match as Sequence")
    match attr:
        case {"some.attribute": a, "other.attribute": b, "dependent": c}:
            print("Matched as Mapping individually", a, b, c)
        case _:
            print("Didn't match a Mapping")
    match attr:
        case {0: a, 1: b}:
            print("Matched as Mapping with 2 int keys", a, b)
        case _:
            print("Didn't match as Mapping with 2 int keys")
print("Registered as Mapping only:")
test(op.attributes)
print("\nAfter additonally registering as Sequence:")
Sequence.register(OpAttributeMap)
test(op.attributes)
```
Output:
```
Registered as Mapping only:
Didn't match as Sequence
Matched as Mapping individually 1 : i8 3.000000e+00 : f64 "text"
Matched as Mapping with 2 int keys NamedAttribute(dependent="text") NamedAttribute(other.attribute=3.000000e+00 : f64)

After additonally registering as Sequence:
matched a Sequence ['dependent', 'other.attribute', 'some.attribute']
Didn't match a Mapping
Didn't match as Mapping with 2 int keys
```
</details>

makslevental Would be great if you could take a look again ❤️

---------

Co-authored-by: Maksim Levental <maksim.levental@gmail.com>
2026-01-10 09:09:05 -08:00
..