[Mlir-commits] [mlir] [mlir:python] Set `__module__` classes generated by nanobind adaptors. (PR #127584)

[Ingo Müller]

https://github.com/ingomueller-net created https://github.com/llvm/llvm-project/pull/127584

This 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.

>From e0d5ad7a7c6cae9b87b568718032978e69acfc2c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ingo=20M=C3=BCller?= <ingomueller at google.com>
Date: Tue, 18 Feb 2025 07:40:45 +0000
Subject: [PATCH] [mlir:python] Set `__module__` classes generated by nanobind
 adaptors.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This 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.

Signed-off-by: Ingo Müller <ingomueller at google.com>
---
 mlir/include/mlir/Bindings/Python/NanobindAdaptors.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 517351cac6dbc..11338b2d49308 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -349,6 +349,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>