[Mlir-commits] [mlir] [MLIR] [Python] Fixed the return type of static_typeid (PR #183021)

[Sergei Lebedev]

https://github.com/superbobry updated https://github.com/llvm/llvm-project/pull/183021

>From dd82ea3c464bd8f620142bfdf3f02db54afa1318 Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev at google.com>
Date: Tue, 24 Feb 2026 09:49:52 +0000
Subject: [PATCH] [MLIR] [Python] Fixed a few issues in the type annotations

* The return type of `static_typeid` in custom attributes is now qualified.
* `DefaultingPyMlir*` helpers also produce qualified types, e.g.
  `_mlir.ir.Location` instead of bare `Location`.
* `loc_tracebacks` uses `Generator` as the return type, since this is what
  `contextmanager` expects in typeshed.

I also sent wjakob/nanobind#1302 and wjakob/nanobind#1303, which should improve
the quality of the auto-generated stubs.
---
 mlir/include/mlir/Bindings/Python/IRCore.h | 36 +++++++++-------------
 mlir/lib/Bindings/Python/IRAttributes.cpp  |  3 +-
 mlir/python/mlir/ir.py                     |  4 +--
 3 files changed, 18 insertions(+), 25 deletions(-)

diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 5953f26d07370..bd2d49acbf681 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -278,7 +278,7 @@ class MLIR_PYTHON_API_EXPORTED DefaultingPyMlirContext
     : public Defaulting<DefaultingPyMlirContext, PyMlirContext> {
 public:
   using Defaulting::Defaulting;
-  static constexpr const char kTypeDescription[] = "Context";
+  static constexpr const char kTypeDescription[] = "_mlir.ir.Context";
   static PyMlirContext &resolve();
 };
 
@@ -524,7 +524,7 @@ class MLIR_PYTHON_API_EXPORTED DefaultingPyLocation
     : public Defaulting<DefaultingPyLocation, PyLocation> {
 public:
   using Defaulting::Defaulting;
-  static constexpr const char kTypeDescription[] = "Location";
+  static constexpr const char kTypeDescription[] = "_mlir.ir.Location";
   static PyLocation &resolve();
 
   operator MlirLocation() const { return *get(); }
@@ -957,16 +957,12 @@ class MLIR_PYTHON_API_EXPORTED PyConcreteType : public BaseTy {
     auto cls = ClassTy(m, DerivedTy::pyClassName, nanobind::is_generic());
     cls.def(nanobind::init<PyType &>(), nanobind::keep_alive<0, 1>(),
             nanobind::arg("cast_from_type"));
-    cls.def_prop_ro_static(
-        "static_typeid",
-        [](nanobind::object & /*class*/) {
-          if (DerivedTy::getTypeIdFunction)
-            return PyTypeID(DerivedTy::getTypeIdFunction());
-          throw nanobind::attribute_error(
-              (DerivedTy::pyClassName + std::string(" has no typeid."))
-                  .c_str());
-        },
-        nanobind::sig("def static_typeid(/) -> TypeID"));
+    cls.def_prop_ro_static("static_typeid", [](nanobind::object & /*class*/) {
+      if (DerivedTy::getTypeIdFunction)
+        return PyTypeID(DerivedTy::getTypeIdFunction());
+      throw nanobind::attribute_error(
+          (DerivedTy::pyClassName + std::string(" has no typeid.")).c_str());
+    });
     cls.def_prop_ro("typeid", [](PyType &self) {
       return nanobind::cast<PyTypeID>(nanobind::cast(self).attr("typeid"));
     });
@@ -1100,16 +1096,12 @@ class MLIR_PYTHON_API_EXPORTED PyConcreteAttribute : public BaseTy {
           return PyType(attr.getContext(), mlirAttributeGetType(attr))
               .maybeDownCast();
         });
-    cls.def_prop_ro_static(
-        "static_typeid",
-        [](nanobind::object & /*class*/) -> PyTypeID {
-          if (DerivedTy::getTypeIdFunction)
-            return PyTypeID(DerivedTy::getTypeIdFunction());
-          throw nanobind::attribute_error(
-              (DerivedTy::pyClassName + std::string(" has no typeid."))
-                  .c_str());
-        },
-        nanobind::sig("def static_typeid(/) -> TypeID"));
+    cls.def_prop_ro_static("static_typeid", [](nanobind::object & /*class*/) {
+      if (DerivedTy::getTypeIdFunction)
+        return PyTypeID(DerivedTy::getTypeIdFunction());
+      throw nanobind::attribute_error(
+          (DerivedTy::pyClassName + std::string(" has no typeid.")).c_str());
+    });
     cls.def_prop_ro("typeid", [](PyAttribute &self) {
       return nanobind::cast<PyTypeID>(nanobind::cast(self).attr("typeid"));
     });
diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 59eb9b8e81cf0..85a662da2483c 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -426,7 +426,8 @@ void PyIntegerAttribute::bindDerived(ClassTy &c) {
       [](nb::object & /*class*/) {
         return PyTypeID(mlirIntegerAttrGetTypeID());
       },
-      nb::sig("def static_typeid(/) -> TypeID"));
+      nb::sig(
+          "def static_typeid(/) -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID")));
 }
 
 nb::object PyIntegerAttribute::toPyInt(PyIntegerAttribute &self) {
diff --git a/mlir/python/mlir/ir.py b/mlir/python/mlir/ir.py
index f4aa2d6b051c8..f84792d4095f4 100644
--- a/mlir/python/mlir/ir.py
+++ b/mlir/python/mlir/ir.py
@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Generator
 from contextlib import contextmanager
 
 from ._mlir_libs._mlir.ir import *
@@ -22,7 +22,7 @@
 
 
 @contextmanager
-def loc_tracebacks(*, max_depth: int | None = None) -> Iterable[None]:
+def loc_tracebacks(*, max_depth: int | None = None) -> Generator[None]:
     """Enables automatic traceback-based locations for MLIR operations.
 
     Operations created within this context will have their location