Skip to content

[mlir:python] Improve mlir_(attribute|type|value)_subclass for nanobinds stubgen #127584

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 19, 2025
Merged

[mlir:python] Improve mlir_(attribute|type|value)_subclass for nanobinds stubgen #127584

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e0d5ad7
[mlir:python] Set `__module__` classes generated by nanobind adaptors.
ingomueller-net Feb 18, 2025
b1d185b
Fix signature of `isinstance` bindings.
ingomueller-net Feb 18, 2025
5991abe
Minor unrelated doc and formatting improvements (NFC).
ingomueller-net Feb 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 21 additions & 8 deletions mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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>
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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));
}
};

Expand Down