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

@@ -1939,3 +1939,31 @@ struct CPython:
 
         self._inc_total_rc()
         return r
+
+    # ===-------------------------------------------------------------------===#
+    # Python Capsule API
+    # ===-------------------------------------------------------------------===#
+    fn PyCapsule_New(
+        mut self,
+        pointer: OpaquePointer,
+        name: StringSlice,
+        destructor: destructor,
+    ) -> PyObjectPtr:
+        """[Reference](
+        https://docs.python.org/3/c-api/capsule.html#c.PyCapsule_New).
+        """
+        return self.lib.call["PyCapsule_New", PyObjectPtr](
+            pointer, name.unsafe_ptr(), destructor
+        )
+
+    fn PyCapsule_GetPointer(
+        mut self,
+        capsule: PyObjectPtr,
+        name: StringSlice,
+    ) -> OpaquePointer:
+        """[Reference](
+        https://docs.python.org/3/c-api/capsule.html#c.PyCapsule_GetPointer).
+        """
+        return self.lib.call["PyCapsule_GetPointer", OpaquePointer](
+            capsule, name.unsafe_ptr()
+        )

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)