[Mlir-commits] [mlir] 22fea18 - [mlir] Better error message in PybindAdaptors.h

[Alex Zinenko]

Author: Alex Zinenko
Date: 2022-02-01T17:49:18+01:00
New Revision: 22fea18e5f4e693376f288b598273352cd4ab99f

URL: https://github.com/llvm/llvm-project/commit/22fea18e5f4e693376f288b598273352cd4ab99f
DIFF: https://github.com/llvm/llvm-project/commit/22fea18e5f4e693376f288b598273352cd4ab99f.diff

LOG: [mlir] Better error message in PybindAdaptors.h

When attempting to cast a pybind11 handle to an MLIR C API object through
capsules, the binding code would attempt to directly access the "_CAPIPtr"
attribute on the object, leading to a rather obscure AttributeError when the
attribute was missing, e.g., on non-MLIR types. Check for its presence and
throw a TypeError instead.

Depends On D117646

Reviewed By: stellaraccident

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

Added: 
    

Modified: 
    mlir/include/mlir/Bindings/Python/PybindAdaptors.h
    mlir/test/python/dialects/python_test.py

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
index 73cc7e4412fbd..9d5a512a4dbe6 100644
--- a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
@@ -42,13 +42,19 @@ struct type_caster<llvm::Optional<T>> : optional_caster<llvm::Optional<T>> {};
 /// an explicit Capsule (which can happen when two C APIs are communicating
 /// directly via Python) or indirectly by querying the MLIR_PYTHON_CAPI_PTR_ATTR
 /// attribute (through which supported MLIR Python API objects export their
-/// contained API pointer as a capsule). This is intended to be used from
-/// type casters, which are invoked with a raw handle (unowned). The returned
-/// object's lifetime may not extend beyond the apiObject handle without
-/// explicitly having its refcount increased (i.e. on return).
+/// contained API pointer as a capsule). Throws a type error if the object is
+/// neither. This is intended to be used from type casters, which are invoked
+/// with a raw handle (unowned). The returned object's lifetime may not extend
+/// beyond the apiObject handle without explicitly having its refcount increased
+/// (i.e. on return).
 static py::object mlirApiObjectToCapsule(py::handle apiObject) {
   if (PyCapsule_CheckExact(apiObject.ptr()))
     return py::reinterpret_borrow<py::object>(apiObject);
+  if (!py::hasattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR)) {
+    auto repr = py::repr(apiObject).cast<std::string>();
+    throw py::type_error(
+        (llvm::Twine("Expected an MLIR object (got ") + repr + ").").str());
+  }
   return apiObject.attr(MLIR_PYTHON_CAPI_PTR_ATTR);
 }
 

diff  --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
index 9c096577f9360..e7b1f44a3ad8e 100644
--- a/mlir/test/python/dialects/python_test.py
+++ b/mlir/test/python/dialects/python_test.py
@@ -247,6 +247,15 @@ def testCustomAttribute():
     else:
       raise
 
+    # The following must trigger a TypeError from our adaptors and must not
+    # crash.
+    try:
+      test.TestAttr(42)
+    except TypeError as e:
+      assert "Expected an MLIR object" in str(e)
+    else:
+      raise
+
     # The following must trigger a TypeError from pybind (therefore, not
     # checking its message) and must not crash.
     try:
@@ -276,6 +285,15 @@ def testCustomType():
     else:
       raise
 
+    # The following must trigger a TypeError from our adaptors and must not
+    # crash.
+    try:
+      test.TestType(42)
+    except TypeError as e:
+      assert "Expected an MLIR object" in str(e)
+    else:
+      raise
+
     # The following must trigger a TypeError from pybind (therefore, not
     # checking its message) and must not crash.
     try: