[Mlir-commits] [mlir] [MLIR] [Python] Fixed a few issues in the type annotations (PR #183021)
Sergei Lebedev
llvmlistbot at llvm.org
Fri Mar 6 05:34:40 PST 2026
https://github.com/superbobry updated https://github.com/llvm/llvm-project/pull/183021
>From d612e6a690a092c1b2c3cf3727e5c92d7b90c09c 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`.
* `ir.*.__enter__` now returns a concrete type instead of `object`, e.g.
`ir.Context.__enter__` returns `Context`.
* `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/lib/Bindings/Python/IRCore.cpp | 12 +++++---
mlir/python/mlir/ir.py | 4 +--
4 files changed, 27 insertions(+), 34 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/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 4d389c656e58d..b8637c57a3f48 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2988,7 +2988,8 @@ void populateIRCore(nb::module_ &m) {
"Returns True if an error was encountered during diagnostic "
"handling.")
.def("__enter__", &PyDiagnosticHandler::contextEnter,
- "Enters the diagnostic handler as a context manager.")
+ "Enters the diagnostic handler as a context manager.",
+ nb::sig("def __enter__(self, /) -> DiagnosticHandler"))
.def("__exit__", &PyDiagnosticHandler::contextExit, "exc_type"_a.none(),
"exc_value"_a.none(), "traceback"_a.none(),
"Exits the diagnostic handler context manager.");
@@ -3034,7 +3035,8 @@ void populateIRCore(nb::module_ &m) {
&PyMlirContext::createFromCapsule,
"Creates a Context from a capsule wrapping MlirContext.")
.def("__enter__", &PyMlirContext::contextEnter,
- "Enters the context as a context manager.")
+ "Enters the context as a context manager.",
+ nb::sig("def __enter__(self, /) -> Context"))
.def("__exit__", &PyMlirContext::contextExit, "exc_type"_a.none(),
"exc_value"_a.none(), "traceback"_a.none(),
"Exits the context manager.")
@@ -3260,7 +3262,8 @@ void populateIRCore(nb::module_ &m) {
.def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyLocation::createFromCapsule,
"Creates a Location from a capsule wrapping MlirLocation.")
.def("__enter__", &PyLocation::contextEnter,
- "Enters the location as a context manager.")
+ "Enters the location as a context manager.",
+ nb::sig("def __enter__(self, /) -> Location"))
.def("__exit__", &PyLocation::contextExit, "exc_type"_a.none(),
"exc_value"_a.none(), "traceback"_a.none(),
"Exits the location context manager.")
@@ -4290,7 +4293,8 @@ void populateIRCore(nb::module_ &m) {
.def(nb::init<PyBlock &>(), "block"_a,
"Inserts after the last operation but still inside the block.")
.def("__enter__", &PyInsertionPoint::contextEnter,
- "Enters the insertion point as a context manager.")
+ "Enters the insertion point as a context manager.",
+ nb::sig("def __enter__(self, /) -> InsertionPoint"))
.def("__exit__", &PyInsertionPoint::contextExit, "exc_type"_a.none(),
"exc_value"_a.none(), "traceback"_a.none(),
"Exits the insertion point context manager.")
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
More information about the Mlir-commits
mailing list