[Mlir-commits] [mlir] [mlir][python] Make the Context/Operation capsule creation methods work as documented. (PR #76010)

[Maksim Levental]

================
@@ -615,18 +615,35 @@ PyMlirContextRef PyMlirContext::forContext(MlirContext context) {
   auto &liveContexts = getLiveContexts();
   auto it = liveContexts.find(context.ptr);
   if (it == liveContexts.end()) {
-    // Create.
-    PyMlirContext *unownedContextWrapper = new PyMlirContext(context);
-    py::object pyRef = py::cast(unownedContextWrapper);
-    assert(pyRef && "cast to py::object failed");
-    liveContexts[context.ptr] = unownedContextWrapper;
-    return PyMlirContextRef(unownedContextWrapper, std::move(pyRef));
+    throw std::runtime_error(
+        "Cannot use a context that is not owned by the Python bindings.");
   }
+
   // Use existing.
   py::object pyRef = py::cast(it->second);
   return PyMlirContextRef(it->second, std::move(pyRef));
 }
 
+PyMlirContextRef PyMlirContext::stealExternalContext(MlirContext context) {
+  py::gil_scoped_acquire acquire;
+  auto &liveContexts = getLiveContexts();
+  auto it = liveContexts.find(context.ptr);
+  if (it != liveContexts.end()) {
+    throw std::runtime_error(
+        "Cannot transfer ownership of the context to Python "
+        "as it is already owned by Python.");
+  }
+
+  PyMlirContext *unownedContextWrapper = new PyMlirContext(context);
+  // Note that the default return value policy on cast is automatic_reference,
+  // which does not take ownership (delete will not be called).
----------------
makslevental wrote:

Just curious: is this specific to `py::cast` and/or a bug in pybind? The docs say about `return_value_policy::automatic`: 

> This policy falls back to the policy return_value_policy::take_ownership when the return value is a pointer... **This is the default policy for `py::class_`-wrapped types**.

and about `return_value_policy::automatic_reference`:

> As above, but use policy `return_value_policy::reference` when the return value is a pointer. This is the default conversion policy for function arguments when calling Python functions manually from C++ code (i.e. via `handle::operator()`) and the casters in `pybind11/stl.h`.

Are we somehow in either of those two cases here?

https://github.com/llvm/llvm-project/pull/76010