Skip to main content

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    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine,
8    class::PyClassImpl,
9    function::{FuncArgs, OptionalArg, PyArithmeticValue, PyComparisonValue},
10    types::{
11        Callable, Comparable, Constructor, Hashable, Initializer, PyComparisonOp, Representable,
12    },
13};
14
15pub use crate::object::PyWeak;
16
17#[derive(FromArgs)]
18#[allow(dead_code)]
19pub struct WeakNewArgs {
20    #[pyarg(positional)]
21    referent: PyObjectRef,
22    #[pyarg(positional, optional)]
23    callback: OptionalArg<PyObjectRef>,
24}
25
26impl PyPayload for PyWeak {
27    #[inline]
28    fn class(ctx: &Context) -> &'static Py<PyType> {
29        ctx.types.weakref_type
30    }
31}
32
33impl Callable for PyWeak {
34    type Args = ();
35
36    #[inline]
37    fn call(zelf: &Py<Self>, _: Self::Args, vm: &VirtualMachine) -> PyResult {
38        Ok(vm.unwrap_or_none(zelf.upgrade()))
39    }
40}
41
42impl Constructor for PyWeak {
43    type Args = WeakNewArgs;
44
45    fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
46        // PyArg_UnpackTuple: only process positional args, ignore kwargs.
47        // Subclass __init__ will handle extra kwargs.
48        let mut positional = args.args.into_iter();
49        let referent = positional
50            .next()
51            .ok_or_else(|| vm.new_arity_type_error("__new__", 1..=2, 0))?;
52        let callback = positional.next().filter(|callback| !vm.is_none(callback));
53        if let Some(_extra) = positional.next() {
54            let got = positional.count() + 3;
55            return Err(vm.new_arity_type_error("__new__", 1..=2, got));
56        }
57        let weak = referent.downgrade_with_typ(callback, cls, vm)?;
58        Ok(weak.into())
59    }
60
61    fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
62        unimplemented!("use slot_new")
63    }
64}
65
66impl Initializer for PyWeak {
67    type Args = WeakNewArgs;
68
69    // weakref_tp_init: accepts args but does nothing (all init done in slot_new)
70    fn init(_zelf: &Py<Self>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<()> {
71        Ok(())
72    }
73}
74
75#[pyclass(
76    with(
77        Callable,
78        Hashable,
79        Comparable,
80        Constructor,
81        Initializer,
82        Representable
83    ),
84    flags(BASETYPE)
85)]
86impl Py<PyWeak> {
87    #[pygetset]
88    fn __callback__(&self, vm: &VirtualMachine) -> PyObjectRef {
89        vm.unwrap_or_none(self.get_callback())
90    }
91
92    #[pyclassmethod]
93    fn __class_getitem__(
94        cls: PyTypeRef,
95        object: PyObjectRef,
96        vm: &VirtualMachine,
97    ) -> PyResult<PyGenericAlias> {
98        PyGenericAlias::from_args(cls, object, vm)
99    }
100}
101
102impl Hashable for PyWeak {
103    fn hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyHash> {
104        let hash = match zelf.hash.load(Ordering::Relaxed) {
105            hash::SENTINEL => {
106                let obj = zelf
107                    .upgrade()
108                    .ok_or_else(|| vm.new_type_error("weak object has gone away"))?;
109                let hash = obj.hash(vm)?;
110                match Radium::compare_exchange(
111                    &zelf.hash,
112                    hash::SENTINEL,
113                    hash::fix_sentinel(hash),
114                    Ordering::Relaxed,
115                    Ordering::Relaxed,
116                ) {
117                    Ok(_) => hash,
118                    Err(prev_stored) => prev_stored,
119                }
120            }
121            hash => hash,
122        };
123        Ok(hash)
124    }
125}
126
127impl Comparable for PyWeak {
128    fn cmp(
129        zelf: &Py<Self>,
130        other: &PyObject,
131        op: PyComparisonOp,
132        vm: &VirtualMachine,
133    ) -> PyResult<PyComparisonValue> {
134        op.eq_only(|| {
135            let other = class_or_notimplemented!(Self, other);
136            let both = zelf.upgrade().zip(other.upgrade());
137            match both {
138                // CPython parity (Objects/weakref.c::weakref_richcompare): use
139                // PyObject_RichCompare on the referents, not the bool variant,
140                // so referent __eq__ runs even when referents share identity.
141                Some((a, b)) => {
142                    let res = a.rich_compare(b, PyComparisonOp::Eq, vm)?;
143                    PyArithmeticValue::from_object(vm, res)
144                        .map(|obj| obj.try_to_bool(vm))
145                        .transpose()
146                }
147                None => Ok(zelf.is(other).into()),
148            }
149        })
150    }
151}
152
153/// `__name__` when that lookup yields a string. A missing name or a non-string
154/// value is left off the repr.
155fn instance_name(obj: &PyObject, vm: &VirtualMachine) -> PyResult<Option<String>> {
156    let found =
157        crate::vm::PyMethod::get_special_ex::<false>(obj, identifier!(vm, __name__), vm, true)?;
158    let Some(crate::vm::PyMethod::Attribute(attr)) = found else {
159        return Ok(None);
160    };
161    Ok(attr
162        .downcast_ref::<crate::builtins::PyStr>()
163        .and_then(|name| name.to_str())
164        .map(str::to_owned))
165}
166
167impl Representable for PyWeak {
168    #[inline]
169    fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
170        let id = zelf.get_id();
171        let Some(obj) = zelf.upgrade() else {
172            return Ok(format!("<weakref at {id:#x}; dead>"));
173        };
174        let type_name = obj.class().fully_qualified_name(vm)?;
175        let obj_id = obj.get_id();
176        let suffix = match instance_name(&obj, vm)? {
177            Some(name) => format!(" ({name})"),
178            None => String::new(),
179        };
180        Ok(format!(
181            "<weakref at {id:#x}; to '{type_name}' at {obj_id:#x}{suffix}>"
182        ))
183    }
184}
185
186pub(crate) fn init(context: &'static Context) {
187    PyWeak::extend_class(context, context.types.weakref_type);
188}