Add weakref support to c-api

bschoenmaeckers · bschoenmaeckers · commit 9edf71dd6017 · 2026-06-01T21:31:04.000+02:00
diff --git 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<Interpreter>> {
diff --git a/crates/capi/src/weakrefobject.rs 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::<PyWeak>() {
+            weak.upgrade()
+        } else if let Some(proxy) = reference.downcast_ref::<PyWeakProxy>() {
+            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::<PyInt>();
+
+            let weak_ref = PyWeakrefReference::new(&object_ty).unwrap();
+            let weak_proxy = PyWeakrefProxy::new(&object_ty).unwrap();
+
+            assert!(weak_ref.is_instance_of::<PyWeakrefReference>());
+            assert!(weak_proxy.is_instance_of::<PyWeakrefProxy>());
+        });
+    }
+
+    #[test]
+    fn new_ref_and_get_ref() {
+        Python::attach(|py| {
+            let object_ty = py.get_type::<PyInt>();
+            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::<PyInt>();
+            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
@@ -42,6 +42,21 @@ impl Constructor for PyWeakProxy {
         Self::Args { referent, callback }: Self::Args,
         vm: &VirtualMachine,
     ) -> PyResult<Self> {
+        let weak = Self::new_weak(referent.as_ref(), callback.into_option(), vm)?;
+        Ok(Self { weak })
+    }
+}
+
+crate::common::static_cell! {
+    static WEAK_SUBCLASS: PyTypeRef;
+}
+
+impl PyWeakProxy {
+    fn new_weak(
+        referent: &PyObject,
+        callback: Option<PyObjectRef>,
+        vm: &VirtualMachine,
+    ) -> PyResult<PyRef<PyWeak>> {
         // 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 +67,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<PyObjectRef>,
+        vm: &VirtualMachine,
+    ) -> PyResult<PyRef<Self>> {
+        let weak = Self::new_weak(referent, callback, vm)?;
+        Ok(Self { weak }.into_ref(&vm.ctx))
+    }
+
+    #[must_use]
+    pub fn get_weak(&self) -> &PyRef<PyWeak> {
+        &self.weak
+    }
 }
 
 #[pyclass(with(
