rustpython_vm/builtins/
weakref.rs

1use super::{PyGenericAlias, PyType, PyTypeRef};
2use crate::common::{
3    atomic::{Ordering, Radium},
4    hash::{self, PyHash},
5};
6use crate::{
7    class::PyClassImpl,
8    function::OptionalArg,
9    types::{Callable, Comparable, Constructor, Hashable, PyComparisonOp, Representable},
10    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine,
11};
12
13pub use crate::object::PyWeak;
14
15#[derive(FromArgs)]
16pub struct WeakNewArgs {
17    #[pyarg(positional)]
18    referent: PyObjectRef,
19    #[pyarg(positional, optional)]
20    callback: OptionalArg<PyObjectRef>,
21}
22
23impl PyPayload for PyWeak {
24    fn class(ctx: &Context) -> &'static Py<PyType> {
25        ctx.types.weakref_type
26    }
27}
28
29impl Callable for PyWeak {
30    type Args = ();
31    #[inline]
32    fn call(zelf: &Py<Self>, _: Self::Args, vm: &VirtualMachine) -> PyResult {
33        Ok(vm.unwrap_or_none(zelf.upgrade()))
34    }
35}
36
37impl Constructor for PyWeak {
38    type Args = WeakNewArgs;
39
40    fn py_new(
41        cls: PyTypeRef,
42        Self::Args { referent, callback }: Self::Args,
43        vm: &VirtualMachine,
44    ) -> PyResult {
45        let weak = referent.downgrade_with_typ(callback.into_option(), cls, vm)?;
46        Ok(weak.into())
47    }
48}
49
50#[pyclass(
51    with(Callable, Hashable, Comparable, Constructor, Representable),
52    flags(BASETYPE)
53)]
54impl PyWeak {
55    #[pyclassmethod(magic)]
56    fn class_getitem(cls: PyTypeRef, args: PyObjectRef, vm: &VirtualMachine) -> PyGenericAlias {
57        PyGenericAlias::new(cls, args, vm)
58    }
59}
60
61impl Hashable for PyWeak {
62    fn hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyHash> {
63        let hash = match zelf.hash.load(Ordering::Relaxed) {
64            hash::SENTINEL => {
65                let obj = zelf
66                    .upgrade()
67                    .ok_or_else(|| vm.new_type_error("weak object has gone away".to_owned()))?;
68                let hash = obj.hash(vm)?;
69                match Radium::compare_exchange(
70                    &zelf.hash,
71                    hash::SENTINEL,
72                    hash::fix_sentinel(hash),
73                    Ordering::Relaxed,
74                    Ordering::Relaxed,
75                ) {
76                    Ok(_) => hash,
77                    Err(prev_stored) => prev_stored,
78                }
79            }
80            hash => hash,
81        };
82        Ok(hash)
83    }
84}
85
86impl Comparable for PyWeak {
87    fn cmp(
88        zelf: &Py<Self>,
89        other: &PyObject,
90        op: PyComparisonOp,
91        vm: &VirtualMachine,
92    ) -> PyResult<crate::function::PyComparisonValue> {
93        op.eq_only(|| {
94            let other = class_or_notimplemented!(Self, other);
95            let both = zelf.upgrade().and_then(|s| other.upgrade().map(|o| (s, o)));
96            let eq = match both {
97                Some((a, b)) => vm.bool_eq(&a, &b)?,
98                None => zelf.is(other),
99            };
100            Ok(eq.into())
101        })
102    }
103}
104
105impl Representable for PyWeak {
106    #[inline]
107    fn repr_str(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
108        let id = zelf.get_id();
109        let repr = if let Some(o) = zelf.upgrade() {
110            format!(
111                "<weakref at {:#x}; to '{}' at {:#x}>",
112                id,
113                o.class().name(),
114                o.get_id(),
115            )
116        } else {
117            format!("<weakref at {id:#x}; dead>")
118        };
119        Ok(repr)
120    }
121}
122
123pub fn init(context: &Context) {
124    PyWeak::extend_class(context, context.types.weakref_type);
125}