From d411a2b9d0cfbc59764ccc32f0396e823a94e0df Mon Sep 17 00:00:00 2001 From: Bas Schoenmaeckers <7943856+bschoenmaeckers@users.noreply.github.com> Date: Tue, 2 Jun 2026 08:31:21 +0200 Subject: [PATCH] Add weakref support to c-api (#8006) --- crates/capi/src/lib.rs | 1 + crates/capi/src/weakrefobject.rs | 106 ++++++++++++++++++++++++++++ crates/vm/src/builtins/weakproxy.rs | 37 ++++++++-- 3 files changed, 137 insertions(+), 7 deletions(-) create mode 100644 crates/capi/src/weakrefobject.rs diff --git a/crates/capi/src/lib.rs b/crates/capi/src/lib.rs index a2bc31b17..2daee1f26 100644 --- a/crates/capi/src/lib.rs +++ b/crates/capi/src/lib.rs @@ -30,6 +30,7 @@ pub mod traceback; pub mod tupleobject; pub mod unicodeobject; mod util; +pub mod weakrefobject; /// Get main interpreter of this process. Will be None if it has not been initialized yet. pub fn get_main_interpreter() -> MutexGuard<'static, Option> { diff --git a/crates/capi/src/weakrefobject.rs b/crates/capi/src/weakrefobject.rs new file mode 100644 index 000000000..706de84b5 --- /dev/null +++ b/crates/capi/src/weakrefobject.rs @@ -0,0 +1,106 @@ +use crate::PyObject; +use crate::object::define_py_check; +use crate::pystate::with_vm; +use core::ffi::c_int; +use rustpython_vm::builtins::{PyWeak, PyWeakProxy}; + +define_py_check!(fn PyWeakref_CheckProxy, types.weakproxy_type); +define_py_check!(fn PyWeakref_CheckRef, types.weakref_type); + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyWeakref_GetRef( + reference: *mut PyObject, + result: *mut *mut PyObject, +) -> c_int { + with_vm(|vm| { + unsafe { + *result = core::ptr::null_mut(); + } + + let reference = unsafe { &*reference }; + let upgraded = if let Some(weak) = reference.downcast_ref::() { + weak.upgrade() + } else if let Some(proxy) = reference.downcast_ref::() { + proxy.get_weak().upgrade() + } else { + return Err(vm.new_type_error("expected a weakref")); + }; + + if let Some(obj) = upgraded { + unsafe { + *result = obj.into_raw().as_ptr(); + } + Ok(true) + } else { + Ok(false) + } + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyWeakref_NewProxy( + ob: *mut PyObject, + callback: *mut PyObject, +) -> *mut PyObject { + with_vm(|vm| { + let ob = unsafe { &*ob }; + let callback = unsafe { callback.as_ref() } + .filter(|callback| !vm.is_none(callback)) + .map(ToOwned::to_owned); + PyWeakProxy::new_weakproxy(ob, callback, vm) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyWeakref_NewRef( + ob: *mut PyObject, + callback: *mut PyObject, +) -> *mut PyObject { + with_vm(|vm| { + let ob = unsafe { &*ob }; + let callback = unsafe { callback.as_ref() } + .filter(|callback| !vm.is_none(callback)) + .map(ToOwned::to_owned); + ob.downgrade(callback, vm) + }) +} + +#[cfg(false)] +mod tests { + use pyo3::prelude::*; + use pyo3::types::PyAnyMethods; + use pyo3::types::{PyInt, PyWeakrefMethods, PyWeakrefProxy, PyWeakrefReference}; + + #[test] + fn check_ref_and_proxy() { + Python::attach(|py| { + let object_ty = py.get_type::(); + + let weak_ref = PyWeakrefReference::new(&object_ty).unwrap(); + let weak_proxy = PyWeakrefProxy::new(&object_ty).unwrap(); + + assert!(weak_ref.is_instance_of::()); + assert!(weak_proxy.is_instance_of::()); + }); + } + + #[test] + fn new_ref_and_get_ref() { + Python::attach(|py| { + let object_ty = py.get_type::(); + let weak_ref = PyWeakrefReference::new(&object_ty).unwrap(); + + assert!(weak_ref.upgrade().is_some_and(|obj| obj.is(&object_ty))); + }); + } + + #[test] + fn new_proxy() { + Python::attach(|py| { + let object_ty = py.get_type::(); + let weak_proxy = PyWeakrefProxy::new(&object_ty).unwrap(); + + assert!(weak_proxy.upgrade().is_some_and(|obj| obj.is(&object_ty))); + }); + } +} diff --git a/crates/vm/src/builtins/weakproxy.rs b/crates/vm/src/builtins/weakproxy.rs index 077e8c349..1bdd7721f 100644 --- a/crates/vm/src/builtins/weakproxy.rs +++ b/crates/vm/src/builtins/weakproxy.rs @@ -42,6 +42,22 @@ impl Constructor for PyWeakProxy { Self::Args { referent, callback }: Self::Args, vm: &VirtualMachine, ) -> PyResult { + let weak = Self::new_weak(referent.as_ref(), callback.into_option(), vm)?; + // TODO: PyWeakProxy should use the same payload as PyWeak + Ok(Self { weak }) + } +} + +crate::common::static_cell! { + static WEAK_SUBCLASS: PyTypeRef; +} + +impl PyWeakProxy { + fn new_weak( + referent: &PyObject, + callback: Option, + vm: &VirtualMachine, + ) -> PyResult> { // using an internal subclass as the class prevents us from getting the generic weakref, // which would mess up the weakref count let weak_cls = WEAK_SUBCLASS.get_or_init(|| { @@ -52,15 +68,22 @@ impl Constructor for PyWeakProxy { super::PyWeak::make_slots(), ) }); - // TODO: PyWeakProxy should use the same payload as PyWeak - Ok(Self { - weak: referent.downgrade_with_typ(callback.into_option(), weak_cls.clone(), vm)?, - }) + referent.downgrade_with_typ(callback, weak_cls.clone(), vm) } -} -crate::common::static_cell! { - static WEAK_SUBCLASS: PyTypeRef; + pub fn new_weakproxy( + referent: &PyObject, + callback: Option, + vm: &VirtualMachine, + ) -> PyResult> { + let weak = Self::new_weak(referent, callback, vm)?; + Ok(Self { weak }.into_ref(&vm.ctx)) + } + + #[must_use] + pub fn get_weak(&self) -> &PyRef { + &self.weak + } } #[pyclass(with(