[mlir][py] avoid crashing on None contexts in custom gets

[ftynse]

Following a series of refactorings, MLIR Python bindings would crash if a
dialect object requiring a context defined using mlir_attribute/type_subclass
was constructed outside of the ir.Context context manager. The type caster
for MlirContext would try using ir.Context.current when the default None
value was provided to the get, which would also just return None. The
caster would then attempt to obtain the MLIR capsule for that None, fail,
but access it anyway without checking, leading to a C++ assertion failure or
segfault.

Guard against this case in nanobind adaptors. Also emit a warning to the user
to clarify expectations, as the default message confusingly says that None is
accepted as context and then fails with a type error. Using Python C API is
currently recommended by nanobind in this case since the surrounding function
must be marked noexcept.

The corresponding test is in the PDL dialect since it is where I first observed
the behavior. Core types are not using the mlir_type_subclass mechanism and
are immune to the problem, so cannot be used for checking.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Oleksandr "Alex" Zinenko (ftynse)

Changes

Following a series of refactorings, MLIR Python bindings would crash if a
dialect object requiring a context defined using mlir_attribute/type_subclass
was constructed outside of the ir.Context context manager. The type caster
for MlirContext would try using ir.Context.current when the default None
value was provided to the get, which would also just return None. The
caster would then attempt to obtain the MLIR capsule for that None, fail,
but access it anyway without checking, leading to a C++ assertion failure or
segfault.

Guard against this case in nanobind adaptors. Also emit a warning to the user
to clarify expectations, as the default message confusingly says that None is
accepted as context and then fails with a type error. Using Python C API is
currently recommended by nanobind in this case since the surrounding function
must be marked noexcept.

The corresponding test is in the PDL dialect since it is where I first observed
the behavior. Core types are not using the mlir_type_subclass mechanism and
are immune to the problem, so cannot be used for checking.

---

Full diff: https://github.com/llvm/llvm-project/pull/171140.diff

2 Files Affected:

• (modified) mlir/include/mlir/Bindings/Python/NanobindAdaptors.h (+14-6)
• (modified) mlir/test/python/dialects/pdl_types.py (+13)

diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 847951ab5fd46..6594670abaaa7 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -181,14 +181,22 @@ struct type_caster<MlirContext> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (src.is_none()) {
       // Gets the current thread-bound context.
-      // TODO: This raises an error of "No current context" currently.
-      // Update the implementation to pretty-print the helpful error that the
-      // core implementations print in this case.
       src = mlir::python::irModule().attr("Context").attr("current");
     }
-    std::optional<nanobind::object> capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToContext(capsule->ptr());
-    return !mlirContextIsNull(value);
+    // If there is no context, including thread-bound, emit a warning (since
+    // this function is not allowed to throw) and fail to cast.
+    if (src.is_none()) {
+      PyErr_Warn(
+          PyExc_RuntimeWarning,
+          "Passing None as MLIR Context is only allowed inside "
+          "the " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " context manager.");
+      return false;
+    }
+    if (std::optional<nanobind::object> capsule = mlirApiObjectToCapsule(src)) {
+      value = mlirPythonCapsuleToContext(capsule->ptr());
+      return !mlirContextIsNull(value);
+    }
+    return false;
   }
 };
 
diff --git a/mlir/test/python/dialects/pdl_types.py b/mlir/test/python/dialects/pdl_types.py
index 16a41e2a4c1ce..bbc6007511721 100644
--- a/mlir/test/python/dialects/pdl_types.py
+++ b/mlir/test/python/dialects/pdl_types.py
@@ -148,3 +148,16 @@ def test_value_type():
     print(parsedType)
     # CHECK: !pdl.value
     print(constructedType)
+
+
+# CHECK-LABEL: TEST: test_type_without_context
+@run
+def test_type_without_context():
+  # Constructing a type without the surrounding ir.Context context manager
+  # should raise an exception but not crash.
+  try:
+    constructedType = pdl.ValueType.get()
+  except TypeError:
+    pass
+  else:
+    assert False, "Expected TypeError to be raised."

[github-actions]

✅ With the latest revision this PR passed the Python code formatter.

[makslevental]

LGTM

[makslevental]

nit: this is the wrong error - it should be RuntimeError - but I guess nanobind is making the choice here not us 🤷

[ftynse]

Yep, its reasoning is that if a type caster fails, it's a type error.