diff --git a/crates/capi/src/lib.rs b/crates/capi/src/lib.rs index 17a5af5ee..90ef71da7 100644 --- a/crates/capi/src/lib.rs +++ b/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; diff --git a/crates/capi/src/modsupport.rs b/crates/capi/src/modsupport.rs new file mode 100644 index 000000000..469a6e639 --- /dev/null +++ b/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::(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::(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::(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::(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::(); + 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::().unwrap(), "hello_mod"); + + let int_val = m.getattr("INT_CONST").unwrap(); + assert_eq!(int_val.extract::().unwrap(), 42); + + let obj_val = m.getattr("OBJ_ATTR").unwrap(); + assert_eq!(obj_val.extract::().unwrap(), "custom_val"); + + let ty_val = m.getattr("int").unwrap(); + assert!(ty_val.is_instance_of::()); + + 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); + } + }); + } +} diff --git a/crates/capi/src/moduleobject.rs b/crates/capi/src/moduleobject.rs index d6a0b72fc..fdd09aeb2 100644 --- a/crates/capi/src/moduleobject.rs +++ b/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::(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::()); + + 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::().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::()); + }); + } +} diff --git a/crates/capi/src/util.rs b/crates/capi/src/util.rs index 495c7f5fd..5d0e2c26f 100644 --- a/crates/capi/src/util.rs +++ b/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>(