[Mlir-commits] [mlir] [MLIR][Python] fixup Context and Location stubs (PR #161433)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/161433

>From fbf833fc43eb1a28a357b2a83a42674b35e28b1e Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Tue, 30 Sep 2025 13:27:55 -0700
Subject: [PATCH] [MLIR][Python] fixup Context and Location stubs

---
 mlir/lib/Bindings/Python/IRCore.cpp  | 11 +++++------
 mlir/lib/Bindings/Python/IRModule.h  |  6 ++----
 mlir/lib/Bindings/Python/IRTypes.cpp |  4 ++--
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 83a8757bb72c7..15a73aa19ef0f 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3597,12 +3597,11 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           },
           "Reports if the operation is attached to its parent block.")
       .def("erase", [](PyOperationBase &self) { self.getOperation().erase(); })
-      .def(
-          "walk", &PyOperationBase::walk, nb::arg("callback"),
-          nb::arg("walk_order") = MlirWalkPostOrder,
-          // clang-format off
-          nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder = " MAKE_MLIR_PYTHON_QUALNAME("ir.WalkOrder.POST_ORDER") ") -> None")
-          // clang-format on
+      .def("walk", &PyOperationBase::walk, nb::arg("callback"),
+           nb::arg("walk_order") = MlirWalkPostOrder,
+           // clang-format off
+          nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder) -> None")
+           // clang-format on
       );
 
   nb::class_<PyOperation, PyOperationBase>(m, "Operation")
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 598ae0188464a..edbd73eade906 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -273,8 +273,7 @@ class DefaultingPyMlirContext
     : public Defaulting<DefaultingPyMlirContext, PyMlirContext> {
 public:
   using Defaulting::Defaulting;
-  static constexpr const char kTypeDescription[] =
-      MAKE_MLIR_PYTHON_QUALNAME("ir.Context");
+  static constexpr const char kTypeDescription[] = "Context";
   static PyMlirContext &resolve();
 };
 
@@ -500,8 +499,7 @@ class DefaultingPyLocation
     : public Defaulting<DefaultingPyLocation, PyLocation> {
 public:
   using Defaulting::Defaulting;
-  static constexpr const char kTypeDescription[] =
-      MAKE_MLIR_PYTHON_QUALNAME("ir.Location");
+  static constexpr const char kTypeDescription[] = "Location";
   static PyLocation &resolve();
 
   operator MlirLocation() const { return *get(); }
diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp
index 3488d92250b45..34c5b8dd86a66 100644
--- a/mlir/lib/Bindings/Python/IRTypes.cpp
+++ b/mlir/lib/Bindings/Python/IRTypes.cpp
@@ -1010,7 +1010,7 @@ class PyTupleType : public PyConcreteType<PyTupleType> {
         },
         nb::arg("elements"), nb::arg("context") = nb::none(),
         // clang-format off
-        nb::sig("def get_tuple(elements: Sequence[Type], context: mlir.ir.Context | None = None) -> TupleType"),
+        nb::sig("def get_tuple(elements: Sequence[Type], context: Context | None = None) -> TupleType"),
         // clang-format on
         "Create a tuple type");
     c.def(
@@ -1070,7 +1070,7 @@ class PyFunctionType : public PyConcreteType<PyFunctionType> {
         },
         nb::arg("inputs"), nb::arg("results"), nb::arg("context") = nb::none(),
         // clang-format off
-        nb::sig("def get(inputs: Sequence[Type], results: Sequence[Type], context: mlir.ir.Context | None = None) -> FunctionType"),
+        nb::sig("def get(inputs: Sequence[Type], results: Sequence[Type], context: Context | None = None) -> FunctionType"),
         // clang-format on
         "Gets a FunctionType from a list of input and result types");
     c.def_prop_ro(