[MLIR][Python] make Sliceable inherit from Sequence

[makslevental]

Generates type stubs like

class RegionSequence(Sequence[Region]):
    def __add__(self, arg: RegionSequence, /) -> list[Region]: ...

    def __iter__(self) -> RegionIterator:
        """Returns an iterator over the regions in the sequence."""

[makslevental]

@ingomueller-net can you try this and see if it fixes you?

[ingomueller-net]

@ingomueller-net can you try this and see if it fixes you?

Yes, it does! So nice! I knew you'd come up with something 😉

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

Generates type stubs like

class RegionSequence(Sequence[Region]):
    def __add__(self, arg: RegionSequence, /) -> list[Region]: ...

    def __iter__(self) -> RegionIterator:
        """Returns an iterator over the regions in the sequence."""

WIP (will polish up if it works for us)

---

Full diff: https://github.com/llvm/llvm-project/pull/170551.diff

2 Files Affected:

• (modified) mlir/lib/Bindings/Python/NanobindUtils.h (+12-1)
• (modified) mlir/test/python/ir/operation.py (+4)

diff --git a/mlir/lib/Bindings/Python/NanobindUtils.h b/mlir/lib/Bindings/Python/NanobindUtils.h
index 658e8ad5330ef..4a78ec96346a2 100644
--- a/mlir/lib/Bindings/Python/NanobindUtils.h
+++ b/mlir/lib/Bindings/Python/NanobindUtils.h
@@ -16,9 +16,11 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/DataTypes.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <string>
+#include <typeinfo>
 #include <variant>
 
 template <>
@@ -344,7 +346,16 @@ class Sliceable {
 
   /// Binds the indexing and length methods in the Python class.
   static void bind(nanobind::module_ &m) {
-    auto clazz = nanobind::class_<Derived>(m, Derived::pyClassName)
+    const std::type_info &elemTy = typeid(ElementTy);
+    PyObject *elemTyInfo = nanobind::detail::nb_type_lookup(&elemTy);
+    assert(elemTyInfo &&
+           "expected nb_type_lookup to succeed for Sliceable elemTy");
+    nanobind::handle elemTyName = nanobind::detail::nb_type_name(elemTyInfo);
+    std::string sig = std::string("class ") + Derived::pyClassName +
+                      "(collections.abc.Sequence[" +
+                      nanobind::cast<std::string>(elemTyName) + "])";
+    auto clazz = nanobind::class_<Derived>(m, Derived::pyClassName,
+                                           nanobind::sig(sig.c_str()))
                      .def("__add__", &Sliceable::dunderAdd);
     Derived::bindDerived(clazz);
 
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index ca99c2a985242..d124c284197b8 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -43,6 +43,10 @@ def testTraverseOpRegionBlockIterators():
     )
     op = module.operation
     assert op.context is ctx
+    # Note, __nb_signature__ stores the fully-qualified signature - the actual type stub emitted is
+    # class RegionSequence(Sequence[Region])
+    # CHECK: class RegionSequence(collections.abc.Sequence[mlir._mlir_libs._mlir.ir.Region])
+    print(RegionSequence.__nb_signature__)
     # Get the block using iterators off of the named collections.
     regions = list(op.regions[:])
     blocks = list(regions[0].blocks)

[makslevental]

There's no way for this to fail - the ElementTy needs to be registered before Sliceable[ElementTy] can be registered...

[ftynse]

Don't we have ElementTy::pyClassName similarly to Derived::pyClassName? Maybe not everywhere, but something to consider. This is probably the first time I see typeid being used in non-toy code so I don't know how well it works across platforms.

[makslevental]

Don't we have ElementTy::pyClassName similarly to Derived::pyClassName?

No:

llvm-project/mlir/lib/Bindings/Python/IRModule.h

Line 759 in 4ada42e

class PyRegion {

This is probably the first time I see typeid being used in non-toy code so I don't know how well it works across platforms.

You realize that all of nanobind/pybind hinges on typeid right?

https://github.com/search?q=repo%3Awjakob%2Fnanobind%20typeid&type=code

Specifically the core impl (the cpp to python mapping) uses std::type_info

https://github.com/wjakob/nanobind/blob/892dd8173dd6fad4ec3a50eaa51863ae45c3205d/src/nb_internals.h#L381

https://github.com/wjakob/nanobind/blob/892dd8173dd6fad4ec3a50eaa51863ae45c3205d/include/nanobind/nb_class.h#L564

Ie that's why RTTI is on for the bindings. So we can be sure it works in all the places the bindings work 🙂.

[ingomueller-net]

LGTM but I am not precluding @ftynse's comments...

[ingomueller-net]

Oh, that's a nice way to test this!