[mlir][python] Make types in register_(dialect|operation) more narrow. (llvm#115307)

This PR makes the `pyClass`/`dialectClass` arguments of the pybind11
functions `register_dialect` and `register_operation` as well as their
return types more narrow, concretely, a `py::type` instead of a
`py::object`. As the name of the arguments indicate, they have to be
called with a type instance (a "class"). The PR also updates the typing
stubs of these functions (in the corresponding `.pyi` file), such that
static type checkers are aware of the changed type. With the previous
typing information, `pyright` raised errors on code generated by
tablegen.

Signed-off-by: Ingo Müller <[email protected]>

Author: ingomueller-net

mlir/lib/Bindings/Python/MainModule.cpp

@@ -58,7 +58,7 @@ PYBIND11_MODULE(_mlir, m) {
   // Registration decorators.
   m.def(
       "register_dialect",
-      [](py::object pyClass) {
+      [](py::type pyClass) {
         std::string dialectNamespace =
             pyClass.attr("DIALECT_NAMESPACE").cast<std::string>();
         PyGlobals::get().registerDialectImpl(dialectNamespace, pyClass);
@@ -68,9 +68,9 @@ PYBIND11_MODULE(_mlir, m) {
       "Class decorator for registering a custom Dialect wrapper");
   m.def(
       "register_operation",
-      [](const py::object &dialectClass, bool replace) -> py::cpp_function {
+      [](const py::type &dialectClass, bool replace) -> py::cpp_function {
         return py::cpp_function(
-            [dialectClass, replace](py::object opClass) -> py::object {
+            [dialectClass, replace](py::type opClass) -> py::type {
               std::string operationName =
                   opClass.attr("OPERATION_NAME").cast<std::string>();
               PyGlobals::get().registerOperationImpl(operationName, opClass,

mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi

@@ -8,5 +8,5 @@ class _Globals:
     def append_dialect_search_prefix(self, module_name: str) -> None: ...
     def _check_dialect_module_loaded(self, dialect_namespace: str) -> bool: ...
 
-def register_dialect(dialect_class: type) -> object: ...
-def register_operation(dialect_class: type, *, replace: bool = ...) -> object: ...
+def register_dialect(dialect_class: type) -> type: ...
+def register_operation(dialect_class: type, *, replace: bool = ...) -> type: ...