[Mlir-commits] [mlir] 6b67aac - [mlir:python] Improve `mlir_(attribute|type|value)_subclass` for `nanobind`s `stubgen` (#127584)

[llvmlistbot at llvm.org]

Author: Ingo Müller
Date: 2025-02-19T08:07:49+01:00
New Revision: 6b67aac31377992465c419b5a296f9a7ba8e7984

URL: https://github.com/llvm/llvm-project/commit/6b67aac31377992465c419b5a296f9a7ba8e7984
DIFF: https://github.com/llvm/llvm-project/commit/6b67aac31377992465c419b5a296f9a7ba8e7984.diff

LOG: [mlir:python] Improve `mlir_(attribute|type|value)_subclass` for `nanobind`s `stubgen` (#127584)

This PR makes several improvements to the stubs that are created by
`mlir_(attribute|type|value)_subclass`.

First, the PR sets the `__module__` attribute of the classes generated
by the nanobind adaptors for attributes, types, and values (via
`mlir_(attribute|type|value)_subclass`). By default, the `__module__`
property is set to `importlib._bootstrap`, which isn't where we want the
new class to live. The new logic sets the property to the name of the
module provided as `scope` instead. This also makes nanobind's `stubgen`
generate stubs for those classes properly, which ignores classes whose
`__module__` does not correspond to the module it is generating stubs
for. This resolves #127518.

Second, the PR overwrites the function signatures generated by `stubgen`
to a format that uses the desired type names (e.g., `mlir.ir.Attribute`
instead of `MlirAttribute`).

Finally, the PR piggy-backs some minor doc and style improvements to
`PythonAdaptors.h`.

---------

Signed-off-by: Ingo Müller <ingomueller at google.com>

Added: 
    

Modified: 
    mlir/include/mlir/Bindings/Python/NanobindAdaptors.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 517351cac6dbc..0608182f00b7e 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -23,8 +23,10 @@
 
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
+// clang-format off
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
+// clang-format on
 #include "llvm/ADT/Twine.h"
 
 // Raw CAPI type casters need to be declared before use, so always include them
@@ -349,6 +351,7 @@ class pure_subclass {
     thisClass = metaclass(derivedClassName, nanobind::make_tuple(superClass),
                           attributes);
     scope.attr(derivedClassName) = thisClass;
+    thisClass.attr("__module__") = scope.attr("__name__");
   }
 
   template <typename Func, typename... Extra>
@@ -434,7 +437,7 @@ class mlir_attribute_subclass : public pure_subclass {
                           const nanobind::object &superCls,
                           GetTypeIDFunctionTy getTypeIDFunction = nullptr)
       : pure_subclass(scope, typeClassName, superCls) {
-    // Casting constructor. Note that it hard, if not impossible, to properly
+    // Casting constructor. Note that it is hard, if not impossible, to properly
     // call chain to parent `__init__` in nanobind due to its special handling
     // for init functions that don't have a fully constructed self-reference,
     // which makes it impossible to forward it to `__init__` of a superclass.
@@ -465,10 +468,13 @@ class mlir_attribute_subclass : public pure_subclass {
     thisClass.attr("__new__") = newCf;
 
     // 'isinstance' method.
+    static const char kIsinstanceSig[] =
+        "def isinstance(other_attribute: " MAKE_MLIR_PYTHON_QUALNAME(
+            "ir") ".Attribute) -> bool";
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirAttribute other) { return isaFunction(other); },
-        nanobind::arg("other_attribute"));
+        nanobind::arg("other_attribute"), nanobind::sig(kIsinstanceSig));
     def("__repr__", [superCls, captureTypeName](nanobind::object self) {
       return nanobind::repr(superCls(self))
           .attr("replace")(superCls.attr("__name__"), captureTypeName);
@@ -512,7 +518,7 @@ class mlir_type_subclass : public pure_subclass {
                      const nanobind::object &superCls,
                      GetTypeIDFunctionTy getTypeIDFunction = nullptr)
       : pure_subclass(scope, typeClassName, superCls) {
-    // Casting constructor. Note that it hard, if not impossible, to properly
+    // Casting constructor. Note that it is hard, if not impossible, to properly
     // call chain to parent `__init__` in nanobind due to its special handling
     // for init functions that don't have a fully constructed self-reference,
     // which makes it impossible to forward it to `__init__` of a superclass.
@@ -542,13 +548,17 @@ class mlir_type_subclass : public pure_subclass {
     thisClass.attr("__new__") = newCf;
 
     // 'isinstance' method.
+    static const char kIsinstanceSig[] =
+        "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME(
+            "ir") ".Type) -> bool";
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirType other) { return isaFunction(other); },
-        nanobind::arg("other_type"));
+        nanobind::arg("other_type"), nanobind::sig(kIsinstanceSig));
     def("__repr__", [superCls, captureTypeName](nanobind::object self) {
-      return nanobind::repr(superCls(self))
-          .attr("replace")(superCls.attr("__name__"), captureTypeName);
+      return nanobind::cast<std::string>(
+          nanobind::repr(superCls(self))
+              .attr("replace")(superCls.attr("__name__"), captureTypeName));
     });
     if (getTypeIDFunction) {
       // 'get_static_typeid' method.
@@ -590,7 +600,7 @@ class mlir_value_subclass : public pure_subclass {
                       IsAFunctionTy isaFunction,
                       const nanobind::object &superCls)
       : pure_subclass(scope, valueClassName, superCls) {
-    // Casting constructor. Note that it hard, if not impossible, to properly
+    // Casting constructor. Note that it is hard, if not impossible, to properly
     // call chain to parent `__init__` in nanobind due to its special handling
     // for init functions that don't have a fully constructed self-reference,
     // which makes it impossible to forward it to `__init__` of a superclass.
@@ -620,10 +630,13 @@ class mlir_value_subclass : public pure_subclass {
     thisClass.attr("__new__") = newCf;
 
     // 'isinstance' method.
+    static const char kIsinstanceSig[] =
+        "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME(
+            "ir") ".Value) -> bool";
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirValue other) { return isaFunction(other); },
-        nanobind::arg("other_value"));
+        nanobind::arg("other_value"), nanobind::sig(kIsinstanceSig));
   }
 };