diff --git a/mojo/stdlib/src/python/_cpython.mojo b/mojo/stdlib/src/python/_cpython.mojo index 54051c83f6..2a9585ce0e 100644 --- a/mojo/stdlib/src/python/_cpython.mojo +++ b/mojo/stdlib/src/python/_cpython.mojo @@ -1925,3 +1925,39 @@ struct CPython: self._inc_total_rc() return r + + # ===-------------------------------------------------------------------===# + # Capsules + # ref: https://docs.python.org/3/c-api/capsule.html + # ===-------------------------------------------------------------------===# + + fn PyCapsule_New( + mut self, + pointer: OpaquePointer, + name: StringSlice, + destructor: destructor, + ) -> PyObjectPtr: + """Create a PyCapsule to communicate to another C extension the C API in `pointer`, identified by `name` and with the custom destructor in `destructor`. + + [Reference](https://docs.python.org/3/c-api/capsule.html#c.PyCapsule_New). + """ + # PyObject *PyCapsule_New(void *pointer, const char *name, PyCapsule_Destructor destructor) + var new_capsule = self.lib.call["PyCapsule_New", PyObjectPtr]( + pointer, name.unsafe_ptr().bitcast[c_char](), destructor + ) + self._inc_total_rc() + return new_capsule + + fn PyCapsule_GetPointer( + mut self, + capsule: PyObjectPtr, + name: StringSlice, + ) -> OpaquePointer: + """Extract the pointer to another C extension from a PyCapsule `capsule` with the given `name`. + + [Reference](https://docs.python.org/3/c-api/capsule.html#c.PyCapsule_GetPointer). + """ + # void *PyCapsule_GetPointer(PyObject *capsule, const char *name) + return self.lib.call["PyCapsule_GetPointer", OpaquePointer]( + capsule, name.unsafe_ptr().bitcast[c_char]() + ) diff --git a/mojo/stdlib/test/python/test_python_cpython.mojo b/mojo/stdlib/test/python/test_python_cpython.mojo index 8b51539845..6c2e6dad78 100644 --- a/mojo/stdlib/test/python/test_python_cpython.mojo +++ b/mojo/stdlib/test/python/test_python_cpython.mojo @@ -14,27 +14,58 @@ # RUN: %mojo %s from python import Python, PythonObject +from python._cpython import PyObjectPtr +from memory import UnsafePointer from testing import assert_equal, assert_false, assert_raises, assert_true def test_PyObject_HasAttrString(mut python: Python): - var Cpython_env = python.impl._cpython + var cpython_env = python.impl.cpython() var the_object = PythonObject(0) - var result = Cpython_env[].PyObject_HasAttrString( + var result = cpython_env.PyObject_HasAttrString( the_object.py_object, "__contains__" ) assert_equal(0, result) the_object = Python.list(1, 2, 3) - result = Cpython_env[].PyObject_HasAttrString( + result = cpython_env.PyObject_HasAttrString( the_object.py_object, "__contains__" ) assert_equal(1, result) _ = the_object +fn destructor(capsule: PyObjectPtr) -> None: + pass + + +def test_PyCapsule(mut python: Python): + var cpython_env = python.impl.cpython() + + # Not a PyCapsule, a NULL pointer is expected. + var the_object = PythonObject(0) + var result = cpython_env.PyCapsule_GetPointer( + the_object.py_object, "some_name" + ) + var expected_none = UnsafePointer[NoneType]() + assert_equal(expected_none, result) + + # Build a capsule. + var capsule_impl = UnsafePointer[UInt64].alloc(1) + var capsule = cpython_env.PyCapsule_New( + capsule_impl.bitcast[NoneType](), "some_name", destructor + ) + var capsule_pointer = cpython_env.PyCapsule_GetPointer(capsule, "some_name") + assert_equal(capsule_impl.bitcast[NoneType](), capsule_pointer) + + # Use a different name. + result = cpython_env.PyCapsule_GetPointer(capsule, "some_other_name") + assert_equal(expected_none, result) + + def main(): # initializing Python instance calls init_python var python = Python() test_PyObject_HasAttrString(python) + test_PyCapsule(python)