Customize exception thrown from mlir.Operation.create() python bindings

The default exception handling isn't very user friendly and does not
point accurately to the issue. Instead we can indicate which of the
operands isn't valid and provide contextual information in the error
message.

Differential Revision: https://reviews.llvm.org/D92710

2 files changed, 52 insertions, 5 deletions

diff --git a/mlir/lib/Bindings/Python/IRModules.cpp b/mlir/lib/Bindings/Python/IRModules.cpp
index 3a80064866c0..39a17d053543 100644
--- a/mlir/lib/Bindings/Python/IRModules.cpp
+++ b/mlir/lib/Bindings/Python/IRModules.cpp
@@ -906,11 +906,31 @@ py::object PyOperation::create(
   if (attributes) {
     mlirAttributes.reserve(attributes->size());
     for (auto &it : *attributes) {
-
-      auto name = it.first.cast<std::string>();
-      auto &attribute = it.second.cast<PyAttribute &>();
-      // TODO: Verify attribute originates from the same context.
-      mlirAttributes.emplace_back(std::move(name), attribute);
+      std::string key;
+      try {
+        key = it.first.cast<std::string>();
+      } catch (py::cast_error &err) {
+        std::string msg = "Invalid attribute key (not a string) when "
+                          "attempting to create the operation \"" +
+                          name + "\" (" + err.what() + ")";
+        throw py::cast_error(msg);
+      }
+      try {
+        auto &attribute = it.second.cast<PyAttribute &>();
+        // TODO: Verify attribute originates from the same context.
+        mlirAttributes.emplace_back(std::move(key), attribute);
+      } catch (py::reference_cast_error &) {
+        // This exception seems thrown when the value is "None".
+        std::string msg =
+            "Found an invalid (`None`?) attribute value for the key \"" + key +
+            "\" when attempting to create the operation \"" + name + "\"";
+        throw py::cast_error(msg);
+      } catch (py::cast_error &err) {
+        std::string msg = "Invalid attribute value for the key \"" + key +
+                          "\" when attempting to create the operation \"" +
+                          name + "\" (" + err.what() + ")";
+        throw py::cast_error(msg);
+      }
     }
   }
   // Unpack/validate successors.
diff --git a/mlir/test/Bindings/Python/ir_operation.py b/mlir/test/Bindings/Python/ir_operation.py
index d23e0b6c0b4e..1f6df8626a0a 100644
--- a/mlir/test/Bindings/Python/ir_operation.py
+++ b/mlir/test/Bindings/Python/ir_operation.py
@@ -551,3 +551,30 @@ def testPrintInvalidOperation():
     # CHECK: "module"() ( {
     # CHECK: }) : () -> ()
 run(testPrintInvalidOperation)
+
+
+# CHECK-LABEL: TEST: testCreateWithInvalidAttributes
+def testCreateWithInvalidAttributes():
+  ctx = Context()
+  with Location.unknown(ctx):
+    try:
+      Operation.create("module", attributes={None:StringAttr.get("name")})
+    except Exception as e:
+      # CHECK: Invalid attribute key (not a string) when attempting to create the operation "module" (Unable to cast Python instance of type <class 'NoneType'> to C++ type
+      print(e)
+    try:
+      Operation.create("module", attributes={42:StringAttr.get("name")})
+    except Exception as e:
+      # CHECK: Invalid attribute key (not a string) when attempting to create the operation "module" (Unable to cast Python instance of type <class 'int'> to C++ type
+      print(e)
+    try:
+      Operation.create("module", attributes={"some_key":ctx})
+    except Exception as e:
+      # CHECK: Invalid attribute value for the key "some_key" when attempting to create the operation "module" (Unable to cast Python instance of type <class '_mlir.ir.Context'> to C++ type 'mlir::python::PyAttribute')
+      print(e)
+    try:
+      Operation.create("module", attributes={"some_key":None})
+    except Exception as e:
+      # CHECK: Found an invalid (`None`?) attribute value for the key "some_key" when attempting to create the operation "module"
+      print(e)
+run(testCreateWithInvalidAttributes)