From 68db47ad6ba222c02ab6c256c0e650e4ad23087a Mon Sep 17 00:00:00 2001
From: Alex Zinenko <git@ozinenko.com>
Date: Mon, 8 Dec 2025 15:41:03 +0100
Subject: [PATCH] [mlir][py] avoid crashing on None contexts in custom `get`s

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.
---
 .../mlir/Bindings/Python/NanobindAdaptors.h   | 20 +++++++++++++------
 mlir/test/python/dialects/pdl_types.py        | 13 ++++++++++++
 2 files changed, 27 insertions(+), 6 deletions(-)

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