1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

//! Implementation in line with the python `weakref` module.
//!
//! See also:
//! - [python weakref module](https://docs.python.org/3/library/weakref.html)
//! - [rust weak struct](https://doc.rust-lang.org/std/rc/struct.Weak.html)
//!
pub(crate) use _weakref::make_module;

#[pymodule]
mod _weakref {
    use crate::{
        builtins::{PyDictRef, PyTypeRef, PyWeak},
        PyObjectRef, PyResult, VirtualMachine,
    };

    #[pyattr(name = "ref")]
    fn ref_(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.types.weakref_type.to_owned()
    }
    #[pyattr]
    fn proxy(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.types.weakproxy_type.to_owned()
    }
    #[pyattr(name = "ReferenceType")]
    fn reference_type(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.types.weakref_type.to_owned()
    }
    #[pyattr(name = "ProxyType")]
    fn proxy_type(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.types.weakproxy_type.to_owned()
    }
    #[pyattr(name = "CallableProxyType")]
    fn callable_proxy_type(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.types.weakproxy_type.to_owned()
    }

    #[pyfunction]
    fn getweakrefcount(obj: PyObjectRef) -> usize {
        obj.weak_count().unwrap_or(0)
    }

    #[pyfunction]
    fn getweakrefs(obj: PyObjectRef) -> Vec<PyObjectRef> {
        match obj.get_weak_references() {
            Some(v) => v.into_iter().map(Into::into).collect(),
            None => vec![],
        }
    }

    #[pyfunction]
    fn _remove_dead_weakref(
        dict: PyDictRef,
        key: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        dict._as_dict_inner()
            .delete_if(vm, &*key, |wr| {
                let wr = wr
                    .payload::<PyWeak>()
                    .ok_or_else(|| vm.new_type_error("not a weakref".to_owned()))?;
                Ok(wr.is_dead())
            })
            .map(drop)
    }
}