[Mlir-commits] [mlir] [MLIR] [Python] Fixed a few issues in the type annotations (PR #183021)

[Sergei Lebedev]

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

>From d67b0409127428decca530121590becb66837d99 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

* Removed an explicit `nb::sig` for `static_typeid`. The inferred type would
  work just fine, and unqualified `TypeID`, which was there previously, only
  really works for core types in the `ir` submodule.
* `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  |  9 ++----
 mlir/python/mlir/ir.py                     |  4 +--
 3 files changed, 19 insertions(+), 30 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..82af663daa405 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -421,12 +421,9 @@ void PyIntegerAttribute::bindDerived(ClassTy &c) {
   c.def_prop_ro("value", toPyInt, "Returns the value of the integer attribute");
   c.def("__int__", toPyInt,
         "Converts the value of the integer attribute to a Python int");
-  c.def_prop_ro_static(
-      "static_typeid",
-      [](nb::object & /*class*/) {
-        return PyTypeID(mlirIntegerAttrGetTypeID());
-      },
-      nb::sig("def static_typeid(/) -> TypeID"));
+  c.def_prop_ro_static("static_typeid", [](nb::object & /*class*/) {
+    return PyTypeID(mlirIntegerAttrGetTypeID());
+  });
 }
 
 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