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

[mehdi_amini]

Summary

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.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

939

This is a lot of boilerplate, happy to take a suggestion on how to improve it!

[mehdi_amini]

Add one more test

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

922

I also don't quite understand why casting None to a PyAttribute requires handling this other exception here, while above for the cast to std::string it does not... (maybe something about std::string itself...)

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.cpp

912

This and the py::reference_cast_error are deeper than I have gone. If you see PybindUtils.h, you'll see that I punted back to catching std::exception and just tightly scoping the body because I had trouble with the corners.

922

I've never hit this before but probably because I usually explicitly check for none().

939

This is the version that I improved, believe it or not. There were prior iterations that were more convoluted with respect to checking the failure conditions in a way that wouldn't leak. But agreed. I have a feeling that this method may turn into a performance hotspot and need reworking anyway, so might make sense to save any major uplift to be based on some measurements. There are just a lot of mismatched types and forms that need to be bridged: not sure there is going to be a really elegant way to do it.

mlir/test/Bindings/Python/ir_operation.py

562

Can you make sure these are ValueError (not Exception -- ValueError is most commonly associated with this kind of mismatch)? You may need to throw them as std::invalid_argument for that to be the case.

[mehdi_amini]

mlir/test/Bindings/Python/ir_operation.py

562

I can, but PyBind issue a RuntimeError for cast conversion: https://github.com/pybind/pybind11/blob/master/include/pybind11/detail/common.h#L724

Can you confirm we should stick with ValueError?

[stellaraccident]

mlir/test/Bindings/Python/ir_operation.py

562

That's annoying/inconsistent with typical usage, but better to be internally self consistent I suppose. Leaving it as is sgtm.

mlir/lib/Bindings/Python/IRModules.cpp

for (PyType *result : *results) {
         throw SetPyError(PyExc_ValueError, "result type cannot be None");
       mlirResults.push_back(*result);
     }
   }
   // Unpack/validate attributes.
   if (attributes) {
     mlirAttributes.reserve(attributes->size());
     for (auto &it : *attributes) {
-
-      auto name = it.first.cast<std::string>();
+      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(name), attribute);
+        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.
   if (successors) {
     llvm::SmallVector<MlirBlock, 4> mlirSuccessors;
     mlirSuccessors.reserve(successors->size());
     for (auto *successor : *successors) {
       // TODO: Verify successor originate from the same context.
       if (!successor)
         throw SetPyError(PyExc_ValueError, "successor block cannot be None");
       mlirSuccessors.push_back(successor->get());
     }
   }
 

mlir/test/Bindings/Python/ir_operation.py

with Location.unknown(ctx):
     # This block does not have a terminator, it may crash the custom printer.
     # Verify that we fallback to the generic printer for safety.
     block = module.regions[0].blocks.append()
     print(module)
     # CHECK: // Verification failed, printing generic form
     # 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)