[mlir:python] Add manual typing annotations to mlir.register_* functions.

[ingomueller-net]

This PR adds a manual typing annotations to the register_operation and register_(type|value)_caster functions in the main mlir module. Since those functions return the result nb::cpp_function, which is of type nb::object, the automatic typing annocations are of the form def f() -> object. This isn't particularly precise and leads to type checking errors when the functions are used. Manually defining the annotation with nb::sig solves the problem.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ingo Müller (ingomueller-net)

Changes

This PR adds a manual typing annotations to the register_operation and register_(type|value)_caster functions in the main mlir module. Since those functions return the result nb::cpp_function, which is of type nb::object, the automatic typing annocations are of the form def f() -> object. This isn't particularly precise and leads to type checking errors when the functions are used. Manually defining the annotation with nb::sig solves the problem.

---

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

1 Files Affected:

• (modified) mlir/lib/Bindings/Python/MainModule.cpp (+8)

diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index a14f09f77d2c3..38259c9d962e5 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -102,6 +102,8 @@ NB_MODULE(_mlir, m) {
               return opClass;
             });
       },
+      nb::sig("def register_operation(dialect_class: type, *, "
+              "replace: bool = False) -> typing.Callable[..., type]"),
       "dialect_class"_a, nb::kw_only(), "replace"_a = false,
       "Produce a class decorator for registering an Operation class as part of "
       "a dialect");
@@ -114,6 +116,9 @@ NB_MODULE(_mlir, m) {
           return typeCaster;
         });
       },
+      nb::sig(
+          "def register_type_caster(typeid: mlir.ir.TypeID, *, "
+          "replace: bool = False) -> typing.Callable[...,  mlir.ir.TypeID]"),
       "typeid"_a, nb::kw_only(), "replace"_a = false,
       "Register a type caster for casting MLIR types to custom user types.");
   m.def(
@@ -126,6 +131,9 @@ NB_MODULE(_mlir, m) {
               return valueCaster;
             });
       },
+      nb::sig(
+          "def register_value_caster(typeid: mlir.ir.TypeID, *, "
+          "replace: bool = False) -> typing.Callable[...,  mlir.ir.TypeID]"),
       "typeid"_a, nb::kw_only(), "replace"_a = false,
       "Register a value caster for casting MLIR values to custom user values.");
 

[ingomueller-net]

Hold on, this doesn't work yet. I'll provide details in a sec. Problems resolved with the latest commit.

[makslevental]

LGTM modulo nit!