[Mlir-commits] [mlir] 2a448da - [mlir][python] Make types in register_(dialect|operation) more narrow. (#115307)

[llvmlistbot at llvm.org]

Author: Ingo Müller
Date: 2024-11-11T09:26:15+01:00
New Revision: 2a448da6e63e2cd3dac63d5398bf121e994d7fc2

URL: https://github.com/llvm/llvm-project/commit/2a448da6e63e2cd3dac63d5398bf121e994d7fc2
DIFF: https://github.com/llvm/llvm-project/commit/2a448da6e63e2cd3dac63d5398bf121e994d7fc2.diff

LOG: [mlir][python] Make types in register_(dialect|operation) more narrow. (#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 <ingomueller at google.com>

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/MainModule.cpp
    mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 8da1ab16a4514b..7c27021902de31 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/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,

diff  --git a/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi b/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi
index 42694747e5f24f..03449b70b7fa38 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi
+++ b/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: ...