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, PyRef, PyResult, VirtualMachine,
8    class::PyClassImpl,
9    function::{FuncArgs, OptionalArg},
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    #[inline]
36    fn call(zelf: &Py<Self>, _: Self::Args, vm: &VirtualMachine) -> PyResult {
37        Ok(vm.unwrap_or_none(zelf.upgrade()))
38    }
39}
40
41impl Constructor for PyWeak {
42    type Args = WeakNewArgs;
43
44    fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
45        // PyArg_UnpackTuple: only process positional args, ignore kwargs.
46        // Subclass __init__ will handle extra kwargs.
47        let mut positional = args.args.into_iter();
48        let referent = positional
49            .next()
50            .ok_or_else(|| vm.new_type_error("__new__ expected at least 1 argument, got 0"))?;
51        let callback = positional.next();
52        if let Some(_extra) = positional.next() {
53            return Err(vm.new_type_error(format!(
54                "__new__ expected at most 2 arguments, got {}",
55                3 + positional.count()
56            )));
57        }
58        let weak = referent.downgrade_with_typ(callback, cls, vm)?;
59        Ok(weak.into())
60    }
61
62    fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
63        unimplemented!("use slot_new")
64    }
65}
66
67impl Initializer for PyWeak {
68    type Args = WeakNewArgs;
69
70    // weakref_tp_init: accepts args but does nothing (all init done in slot_new)
71    fn init(_zelf: PyRef<Self>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<()> {
72        Ok(())
73    }
74}
75
76#[pyclass(
77    with(
78        Callable,
79        Hashable,
80        Comparable,
81        Constructor,
82        Initializer,
83        Representable
84    ),
85    flags(BASETYPE)
86)]
87impl PyWeak {
88    #[pyclassmethod]
89    fn __class_getitem__(cls: PyTypeRef, args: PyObjectRef, vm: &VirtualMachine) -> PyGenericAlias {
90        PyGenericAlias::from_args(cls, args, vm)
91    }
92}
93
94impl Hashable for PyWeak {
95    fn hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyHash> {
96        let hash = match zelf.hash.load(Ordering::Relaxed) {
97            hash::SENTINEL => {
98                let obj = zelf
99                    .upgrade()
100                    .ok_or_else(|| vm.new_type_error("weak object has gone away"))?;
101                let hash = obj.hash(vm)?;
102                match Radium::compare_exchange(
103                    &zelf.hash,
104                    hash::SENTINEL,
105                    hash::fix_sentinel(hash),
106                    Ordering::Relaxed,
107                    Ordering::Relaxed,
108                ) {
109                    Ok(_) => hash,
110                    Err(prev_stored) => prev_stored,
111                }
112            }
113            hash => hash,
114        };
115        Ok(hash)
116    }
117}
118
119impl Comparable for PyWeak {
120    fn cmp(
121        zelf: &Py<Self>,
122        other: &PyObject,
123        op: PyComparisonOp,
124        vm: &VirtualMachine,
125    ) -> PyResult<crate::function::PyComparisonValue> {
126        op.eq_only(|| {
127            let other = class_or_notimplemented!(Self, other);
128            let both = zelf.upgrade().and_then(|s| other.upgrade().map(|o| (s, o)));
129            let eq = match both {
130                Some((a, b)) => vm.bool_eq(&a, &b)?,
131                None => zelf.is(other),
132            };
133            Ok(eq.into())
134        })
135    }
136}
137
138impl Representable for PyWeak {
139    #[inline]
140    fn repr_str(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
141        let id = zelf.get_id();
142        let repr = if let Some(o) = zelf.upgrade() {
143            format!(
144                "<weakref at {:#x}; to '{}' at {:#x}>",
145                id,
146                o.class().name(),
147                o.get_id(),
148            )
149        } else {
150            format!("<weakref at {id:#x}; dead>")
151        };
152        Ok(repr)
153    }
154}
155
156pub fn init(context: &'static Context) {
157    PyWeak::extend_class(context, context.types.weakref_type);
158}