[mlir][py] avoid crashing on None contexts in custom gets (#171140)

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.

Author: ftynse

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;
   }
 };
 

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."