feat(capi): implement PyModule_Add* and module helpers from modsupport.h (#8927)

Assisted-by: Gemini:gemini-3.7-flash

crates/capi/src/lib.rs

@@ -27,6 +27,7 @@ pub mod listobject;
 pub mod longobject;
 pub mod memoryobject;
 pub mod methodobject;
+pub mod modsupport;
 pub mod moduleobject;
 pub mod object;
 pub mod objimpl;

crates/capi/src/modsupport.rs

@@ -0,0 +1,148 @@
+use crate::PyObject;
+use crate::object::PyTypeObject;
+use crate::pystate::with_vm;
+use crate::util::{CStrExt, FfiPtrExt};
+use core::ffi::{c_char, c_int, c_long};
+use rustpython_vm::AsObject;
+use rustpython_vm::builtins::PyModule;
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyModule_AddObjectRef(
+    module: *mut PyObject,
+    name: *const c_char,
+    value: *mut PyObject,
+) -> c_int {
+    with_vm(|vm| {
+        let value_obj = unsafe { value.assume_borrowed_or_err(vm) }?;
+        let module = unsafe { module.assume_borrowed_and_cast::<PyModule>(vm) }?;
+        let name_str = unsafe { name.try_as_str(vm) }?;
+        module.dict().set_item(name_str, value_obj.to_owned(), vm)?;
+        Ok(())
+    })
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyModule_AddObject(
+    module: *mut PyObject,
+    name: *const c_char,
+    value: *mut PyObject,
+) -> c_int {
+    let ret = unsafe { PyModule_AddObjectRef(module, name, value) };
+    if ret == 0 && !value.is_null() {
+        let _ = unsafe { value.assume_owned_or_opt() };
+    }
+    ret
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyModule_AddIntConstant(
+    module: *mut PyObject,
+    name: *const c_char,
+    value: c_long,
+) -> c_int {
+    with_vm(|vm| {
+        let module = unsafe { module.assume_borrowed_and_cast::<PyModule>(vm) }?;
+        let name_str = unsafe { name.try_as_str(vm) }?;
+        let py_int = vm.ctx.new_int(value);
+        module.dict().set_item(name_str, py_int.into(), vm)?;
+        Ok(())
+    })
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyModule_AddStringConstant(
+    module: *mut PyObject,
+    name: *const c_char,
+    value: *const c_char,
+) -> c_int {
+    with_vm(|vm| {
+        let module = unsafe { module.assume_borrowed_and_cast::<PyModule>(vm) }?;
+        let name_str = unsafe { name.try_as_str(vm) }?;
+        let val_str = unsafe { value.try_as_str(vm) }?;
+        let py_str = vm.ctx.new_str(val_str);
+        module.dict().set_item(name_str, py_str.into(), vm)?;
+        Ok(())
+    })
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyModule_AddType(
+    module: *mut PyObject,
+    type_: *mut PyTypeObject,
+) -> c_int {
+    with_vm(|vm| {
+        let module = unsafe { module.assume_borrowed_and_cast::<PyModule>(vm) }?;
+        let ty = unsafe { type_.assume_borrowed() };
+        let name = ty.__name__(vm);
+        let slot_name = ty.slot_name();
+        let name_str = name.to_str().unwrap_or(&slot_name);
+        let short_name = name_str.rsplit('.').next().unwrap_or(name_str);
+        module
+            .dict()
+            .set_item(short_name, ty.as_object().to_owned(), vm)?;
+        Ok(())
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use pyo3::prelude::*;
+    use pyo3::types::{PyAnyMethods, PyInt, PyModule, PyString, PyType};
+
+    #[test]
+    fn module_add_constants_and_objects() {
+        Python::attach(|py| {
+            let m = PyModule::new(py, "testmod").unwrap();
+            let c_str_name = c"STR_CONST";
+            let c_str_val = c"hello_mod";
+            let c_int_name = c"INT_CONST";
+            let c_obj_name = c"OBJ_ATTR";
+
+            unsafe {
+                let ret_str = super::PyModule_AddStringConstant(
+                    m.as_ptr().cast(),
+                    c_str_name.as_ptr(),
+                    c_str_val.as_ptr(),
+                );
+                assert_eq!(ret_str, 0);
+
+                let ret_int =
+                    super::PyModule_AddIntConstant(m.as_ptr().cast(), c_int_name.as_ptr(), 42);
+                assert_eq!(ret_int, 0);
+
+                let obj = PyString::new(py, "custom_val");
+                let ret_obj = super::PyModule_AddObjectRef(
+                    m.as_ptr().cast(),
+                    c_obj_name.as_ptr(),
+                    obj.as_ptr().cast(),
+                );
+                assert_eq!(ret_obj, 0);
+
+                let type_ty = py.get_type::<PyInt>();
+                let ret_type = super::PyModule_AddType(m.as_ptr().cast(), type_ty.as_ptr().cast());
+                assert_eq!(ret_type, 0);
+            }
+
+            let str_val = m.getattr("STR_CONST").unwrap();
+            assert_eq!(str_val.extract::<String>().unwrap(), "hello_mod");
+
+            let int_val = m.getattr("INT_CONST").unwrap();
+            assert_eq!(int_val.extract::<i32>().unwrap(), 42);
+
+            let obj_val = m.getattr("OBJ_ATTR").unwrap();
+            assert_eq!(obj_val.extract::<String>().unwrap(), "custom_val");
+
+            let ty_val = m.getattr("int").unwrap();
+            assert!(ty_val.is_instance_of::<PyType>());
+
+            unsafe {
+                let ret_null = super::PyModule_AddObjectRef(
+                    m.as_ptr().cast(),
+                    c"NULL_ATTR".as_ptr(),
+                    core::ptr::null_mut(),
+                );
+                assert_eq!(ret_null, -1);
+            }
+        });
+    }
+}

crates/capi/src/moduleobject.rs

@@ -1,7 +1,8 @@
 use crate::PyObject;
 use crate::object::define_py_check;
 use crate::pystate::with_vm;
-use crate::util::FfiPtrExt;
+use crate::util::{CStrExt, FfiPtrExt};
+use rustpython_vm::AsObject;
 use rustpython_vm::builtins::{PyModule, PyStr};
 
 define_py_check!(fn PyModule_Check, types.module_type);
@@ -41,3 +42,49 @@ pub unsafe extern "C" fn PyModule_NewObject(name: *mut PyObject) -> *mut PyObjec
         Ok(vm.new_module(name, vm.ctx.new_dict(), None))
     })
 }
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyModule_New(name: *const core::ffi::c_char) -> *mut PyObject {
+    with_vm(|vm| {
+        let name_str = unsafe { name.try_as_str(vm) }?;
+        Ok(vm.new_module(name_str, vm.ctx.new_dict(), None))
+    })
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyModule_GetDict(module: *mut PyObject) -> *mut PyObject {
+    with_vm(|vm| {
+        let module = unsafe { module.assume_borrowed_and_cast::<PyModule>(vm) }?;
+        Ok(module.dict().as_object().as_raw().cast_mut())
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use pyo3::prelude::*;
+    use pyo3::types::{PyAnyMethods, PyDict, PyModule, PyString};
+
+    #[test]
+    fn module_new_and_get_dict() {
+        Python::attach(|py| {
+            let mod_name = c"dynamic_mod";
+            let m_ptr = unsafe { super::PyModule_New(mod_name.as_ptr()) };
+            assert!(!m_ptr.is_null());
+            let m = unsafe { pyo3::Bound::from_owned_ptr(py, m_ptr.cast()) };
+            assert!(m.is_instance_of::<PyModule>());
+
+            let name_ptr = unsafe { super::PyModule_GetNameObject(m_ptr) };
+            assert!(!name_ptr.is_null());
+            let name = unsafe { pyo3::Bound::from_owned_ptr(py, name_ptr.cast()) };
+            assert_eq!(
+                name.cast::<PyString>().unwrap().to_str().unwrap(),
+                "dynamic_mod"
+            );
+
+            let dict_ptr = unsafe { super::PyModule_GetDict(m_ptr) };
+            assert!(!dict_ptr.is_null());
+            let dict = unsafe { pyo3::Bound::from_borrowed_ptr(py, dict_ptr.cast()) };
+            assert!(dict.is_instance_of::<PyDict>());
+        });
+    }
+}

crates/capi/src/util.rs

@@ -243,6 +243,16 @@ pub(crate) trait FfiPtrExt: Sized {
     }
     unsafe fn assume_owned(self) -> Self::Owned;
     unsafe fn assume_borrowed_or_opt<'a>(self) -> Option<&'a Self::Borrowed>;
+    unsafe fn assume_borrowed_or_err<'a>(
+        self,
+        vm: &VirtualMachine,
+    ) -> PyResult<&'a Self::Borrowed> {
+        unsafe { self.assume_borrowed_or_opt() }.ok_or_else(|| {
+            vm.take_raised_exception().unwrap_or_else(|| {
+                vm.new_system_error("Native function returned NULL, but there was no exception set")
+            })
+        })
+    }
     unsafe fn assume_borrowed<'a>(self) -> &'a Self::Borrowed;
 
     unsafe fn assume_borrowed_and_cast<'a, T: PyPayload>(