Add PyCapsule_New and PyCapsule_GetPointer

These are critical when interacting with PyCapsules created by other
libraries in the system. The current use case is interacting with
PyArrow.

Signed-off-by: Marius Seritan <[email protected]>

Author: winding-lines

mojo/stdlib/src/python/_cpython.mojo

@@ -1929,3 +1929,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]()
+        )

mojo/stdlib/test/python/test_python_cpython.mojo

@@ -14,6 +14,8 @@
 # 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
 
 
@@ -34,7 +36,34 @@ def test_PyObject_HasAttrString(mut python: Python):
     _ = 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 = UnsafePointer[NoneType]()
+    assert_equal(expected, 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)
+
+
 def main():
     # initializing Python instance calls init_python
     var python = Python()
     test_PyObject_HasAttrString(python)
+    test_PyCapsule(python)