[MLIR][Python] fixup Context and Location stubs and NanobindAdaptors

[makslevental]

add correct names for NB_TYPE_CASTER(..., name) so users of NanobindAdaptors.h can generate the correct hints. Also fix a few straggler stubs.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/161433.diff

3 Files Affected:

• (modified) mlir/lib/Bindings/Python/IRCore.cpp (+1-1)
• (modified) mlir/lib/Bindings/Python/IRModule.h (+2-4)
• (modified) mlir/lib/Bindings/Python/IRTypes.cpp (+2-2)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 83a8757bb72c7..430f59629668a 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3601,7 +3601,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           "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")
+          nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder) -> None")
           // clang-format on
       );
 
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(

[makslevental]

@superbobry can you check this patch with whatever typechecker you were running (I think?) and see if it resolves the circular import issue.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[makslevental]

i removed the default val here because stubgen gets confused (for whatever reason) and tries to do import WalkOrder even though WalkOrder is defined at this point in the stub........... 🤷