Skip to main content

rustpython_vm/builtins/
weakproxy.rs

1use super::{PyStr, PyStrRef, PyType, PyWeak};
2use crate::common::lock::LazyLock;
3use crate::{
4    Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func,
5    class::PyClassImpl,
6    common::hash::PyHash,
7    function::{FuncArgs, PyArithmeticValue, PyComparisonValue, PySetterValue},
8    protocol::{PyIter, PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
9    stdlib::builtins::reversed,
10    types::{
11        AsMapping, AsNumber, AsSequence, Callable, Comparable, Constructor, GetAttr, Hashable,
12        IterNext, Iterable, PyComparisonOp, Representable, SetAttr,
13    },
14};
15
16#[pyclass(module = "weakref", name = "weakproxy", unhashable = true)]
17#[derive(Debug)]
18#[repr(transparent)]
19pub struct PyWeakProxy(PyWeak);
20
21impl PyPayload for PyWeakProxy {
22    const PAYLOAD_TYPE_ID: core::any::TypeId = <PyWeak as PyPayload>::PAYLOAD_TYPE_ID;
23
24    #[inline]
25    unsafe fn validate_downcastable_from(obj: &PyObject) -> bool {
26        <Self as ::rustpython_vm::class::PyClassDef>::BASICSIZE <= obj.class().slots().basicsize
27            && obj
28                .class()
29                .fast_issubclass(<Self as ::rustpython_vm::class::StaticType>::static_type())
30    }
31
32    #[inline]
33    fn class(ctx: &Context) -> &'static Py<PyType> {
34        ctx.types.weakproxy_type
35    }
36}
37
38#[derive(FromArgs)]
39pub struct WeakProxyNewArgs {
40    #[pyarg(positional)]
41    object: PyObjectRef,
42    #[pyarg(positional, optional)]
43    callback: Option<PyObjectRef>,
44}
45
46impl Constructor for PyWeakProxy {
47    type Args = ();
48
49    fn py_new(cls: &Py<PyType>, _args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
50        Err(vm.new_type_error(format!("cannot create '{}' instances", cls.slot_name())))
51    }
52}
53
54impl PyWeakProxy {
55    pub fn from_new_args(args: WeakProxyNewArgs, vm: &VirtualMachine) -> PyResult<PyRef<PyWeak>> {
56        let WeakProxyNewArgs { object, callback } = args;
57        let callback = callback.filter(|callback| !vm.is_none(callback));
58        Self::new_weakproxy(object.as_ref(), callback, vm)
59    }
60
61    pub fn new_weakproxy(
62        referent: &PyObject,
63        callback: Option<PyObjectRef>,
64        vm: &VirtualMachine,
65    ) -> PyResult<PyRef<PyWeak>> {
66        let typ = if referent.is_callable() {
67            vm.ctx.types.weakcallableproxy_type.to_owned()
68        } else {
69            vm.ctx.types.weakproxy_type.to_owned()
70        };
71        referent.downgrade_with_typ(callback, typ, vm)
72    }
73
74    #[must_use]
75    pub fn get_weak(&self) -> &PyWeak {
76        &self.0
77    }
78}
79
80#[pyclass(
81    module = "weakref",
82    name = "weakcallableproxy",
83    base = PyWeakProxy,
84    ctx = "weakcallableproxy_type",
85    unhashable = true
86)]
87#[derive(Debug)]
88#[repr(transparent)]
89pub struct PyWeakCallableProxy(PyWeakProxy);
90
91#[pyclass(with(Callable, Constructor))]
92impl PyWeakCallableProxy {}
93
94impl Constructor for PyWeakCallableProxy {
95    type Args = ();
96
97    fn py_new(cls: &Py<PyType>, _args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
98        Err(vm.new_type_error(format!("cannot create '{}' instances", cls.slot_name())))
99    }
100}
101
102impl Callable for PyWeakCallableProxy {
103    type Args = FuncArgs;
104
105    fn call(zelf: &Py<Self>, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
106        zelf.0.try_upgrade(vm)?.call(args, vm)
107    }
108}
109
110impl PyWeakProxy {
111    fn try_upgrade(&self, vm: &VirtualMachine) -> PyResult {
112        self.0.upgrade().ok_or_else(|| new_reference_error(vm))
113    }
114
115    fn len(&self, vm: &VirtualMachine) -> PyResult<usize> {
116        self.try_upgrade(vm)?.length(vm)
117    }
118    fn __contains__(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
119        self.try_upgrade(vm)?
120            .sequence_unchecked()
121            .contains(needle, vm)
122    }
123
124    fn getitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult {
125        let obj = self.try_upgrade(vm)?;
126        obj.get_item(needle, vm)
127    }
128
129    fn setitem(&self, needle: &PyObject, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
130        let obj = self.try_upgrade(vm)?;
131        obj.set_item(needle, value, vm)
132    }
133
134    fn delitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
135        let obj = self.try_upgrade(vm)?;
136        obj.del_item(needle, vm)
137    }
138}
139
140#[pyclass(with(
141    GetAttr,
142    SetAttr,
143    Constructor,
144    Comparable,
145    AsNumber,
146    AsSequence,
147    AsMapping,
148    Representable,
149    IterNext
150))]
151impl Py<PyWeakProxy> {
152    /// Return str(self).
153    #[pymethod]
154    fn __str__(zelf: &Self, vm: &VirtualMachine) -> PyResult<PyStrRef> {
155        zelf.try_upgrade(vm)?.str(vm)
156    }
157
158    #[pymethod]
159    fn __bytes__(&self, vm: &VirtualMachine) -> PyResult {
160        self.try_upgrade(vm)?.bytes(vm)
161    }
162
163    #[pymethod]
164    fn __reversed__(&self, vm: &VirtualMachine) -> PyResult {
165        let obj = self.try_upgrade(vm)?;
166        reversed(obj, vm)
167    }
168}
169
170impl Iterable for PyWeakProxy {
171    fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
172        let obj = zelf.try_upgrade(vm)?;
173        Ok(obj.get_iter(vm)?.into())
174    }
175}
176
177impl IterNext for PyWeakProxy {
178    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
179        let obj = zelf.try_upgrade(vm)?;
180        if obj.class().slots().iternext.load().is_none() {
181            return Err(vm.new_type_error("Weakref proxy referenced a non-iterator"));
182        }
183        PyIter::new(obj).next(vm)
184    }
185}
186
187fn new_reference_error(vm: &VirtualMachine) -> PyRef<super::PyBaseException> {
188    vm.new_exception_msg(
189        vm.ctx.exceptions.reference_error.to_owned(),
190        "weakly-referenced object no longer exists".into(),
191    )
192}
193
194impl GetAttr for PyWeakProxy {
195    // TODO: callbacks
196    fn getattro(zelf: &Py<Self>, name: &Py<PyStr>, vm: &VirtualMachine) -> PyResult {
197        let obj = zelf.try_upgrade(vm)?;
198        obj.get_attr(name, vm)
199    }
200}
201
202impl SetAttr for PyWeakProxy {
203    fn setattro(
204        zelf: &Py<Self>,
205        attr_name: &Py<PyStr>,
206        value: PySetterValue,
207        vm: &VirtualMachine,
208    ) -> PyResult<()> {
209        let obj = zelf.try_upgrade(vm)?;
210        obj.call_set_attr(vm, attr_name, value)
211    }
212}
213
214fn proxy_upgrade(obj: &PyObject, vm: &VirtualMachine) -> PyResult {
215    obj.downcast_ref::<PyWeakProxy>()
216        .expect("proxy_upgrade called on non-PyWeakProxy object")
217        .try_upgrade(vm)
218}
219
220fn proxy_upgrade_opt(obj: &PyObject, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
221    match obj.downcast_ref::<PyWeakProxy>() {
222        Some(proxy) => Ok(Some(proxy.try_upgrade(vm)?)),
223        None => Ok(None),
224    }
225}
226
227fn proxy_unary_op(
228    obj: &PyObject,
229    vm: &VirtualMachine,
230    op: fn(&VirtualMachine, &PyObject) -> PyResult,
231) -> PyResult {
232    let upgraded = proxy_upgrade(obj, vm)?;
233    op(vm, &upgraded)
234}
235
236macro_rules! proxy_unary_slot {
237    ($vm_method:ident) => {
238        Some(|number, vm| proxy_unary_op(number.obj, vm, |vm, obj| vm.$vm_method(obj)))
239    };
240}
241
242fn proxy_binary_op(
243    a: &PyObject,
244    b: &PyObject,
245    vm: &VirtualMachine,
246    op: fn(&VirtualMachine, &PyObject, &PyObject) -> PyResult,
247) -> PyResult {
248    let a_up = proxy_upgrade_opt(a, vm)?;
249    let b_up = proxy_upgrade_opt(b, vm)?;
250    let a_ref = a_up.as_deref().unwrap_or(a);
251    let b_ref = b_up.as_deref().unwrap_or(b);
252    op(vm, a_ref, b_ref)
253}
254
255macro_rules! proxy_binary_slot {
256    ($vm_method:ident) => {
257        Some(|a, b, vm| proxy_binary_op(a, b, vm, |vm, a, b| vm.$vm_method(a, b)))
258    };
259}
260
261fn proxy_ternary_op(
262    a: &PyObject,
263    b: &PyObject,
264    c: &PyObject,
265    vm: &VirtualMachine,
266    op: fn(&VirtualMachine, &PyObject, &PyObject, &PyObject) -> PyResult,
267) -> PyResult {
268    let a_up = proxy_upgrade_opt(a, vm)?;
269    let b_up = proxy_upgrade_opt(b, vm)?;
270    let c_up = proxy_upgrade_opt(c, vm)?;
271    let a_ref = a_up.as_deref().unwrap_or(a);
272    let b_ref = b_up.as_deref().unwrap_or(b);
273    let c_ref = c_up.as_deref().unwrap_or(c);
274    op(vm, a_ref, b_ref, c_ref)
275}
276
277macro_rules! proxy_ternary_slot {
278    ($vm_method:ident) => {
279        Some(|a, b, c, vm| proxy_ternary_op(a, b, c, vm, |vm, a, b, c| vm.$vm_method(a, b, c)))
280    };
281}
282
283impl AsNumber for PyWeakProxy {
284    fn as_number() -> &'static PyNumberMethods {
285        static AS_NUMBER: LazyLock<PyNumberMethods> = LazyLock::new(|| PyNumberMethods {
286            boolean: Some(|number, vm| {
287                let obj = proxy_upgrade(number.obj, vm)?;
288                obj.is_true(vm)
289            }),
290            int: Some(|number, vm| {
291                let obj = proxy_upgrade(number.obj, vm)?;
292                obj.try_int(vm).map(Into::into)
293            }),
294            float: Some(|number, vm| {
295                let obj = proxy_upgrade(number.obj, vm)?;
296                obj.try_float(vm).map(Into::into)
297            }),
298            index: Some(|number, vm| {
299                let obj = proxy_upgrade(number.obj, vm)?;
300                obj.try_index(vm).map(Into::into)
301            }),
302            negative: proxy_unary_slot!(_neg),
303            positive: proxy_unary_slot!(_pos),
304            absolute: proxy_unary_slot!(_abs),
305            invert: proxy_unary_slot!(_invert),
306            add: proxy_binary_slot!(_add),
307            subtract: proxy_binary_slot!(_sub),
308            multiply: proxy_binary_slot!(_mul),
309            remainder: proxy_binary_slot!(_mod),
310            divmod: proxy_binary_slot!(_divmod),
311            lshift: proxy_binary_slot!(_lshift),
312            rshift: proxy_binary_slot!(_rshift),
313            and: proxy_binary_slot!(_and),
314            xor: proxy_binary_slot!(_xor),
315            or: proxy_binary_slot!(_or),
316            floor_divide: proxy_binary_slot!(_floordiv),
317            true_divide: proxy_binary_slot!(_truediv),
318            matrix_multiply: proxy_binary_slot!(_matmul),
319            inplace_add: proxy_binary_slot!(_iadd),
320            inplace_subtract: proxy_binary_slot!(_isub),
321            inplace_multiply: proxy_binary_slot!(_imul),
322            inplace_remainder: proxy_binary_slot!(_imod),
323            inplace_lshift: proxy_binary_slot!(_ilshift),
324            inplace_rshift: proxy_binary_slot!(_irshift),
325            inplace_and: proxy_binary_slot!(_iand),
326            inplace_xor: proxy_binary_slot!(_ixor),
327            inplace_or: proxy_binary_slot!(_ior),
328            inplace_floor_divide: proxy_binary_slot!(_ifloordiv),
329            inplace_true_divide: proxy_binary_slot!(_itruediv),
330            inplace_matrix_multiply: proxy_binary_slot!(_imatmul),
331            power: proxy_ternary_slot!(_pow),
332            inplace_power: proxy_ternary_slot!(_ipow),
333        });
334        &AS_NUMBER
335    }
336}
337
338impl Comparable for PyWeakProxy {
339    fn cmp(
340        zelf: &Py<Self>,
341        other: &PyObject,
342        op: PyComparisonOp,
343        vm: &VirtualMachine,
344    ) -> PyResult<PyComparisonValue> {
345        let obj = zelf.try_upgrade(vm)?;
346        // CPython parity (Objects/weakref.c::proxy_richcompare): delegate to
347        // PyObject_RichCompare on the referent, not the bool variant.
348        let res = obj.rich_compare(other.to_owned(), op, vm)?;
349        PyArithmeticValue::from_object(vm, res)
350            .map(|o| o.try_to_bool(vm))
351            .transpose()
352    }
353}
354
355impl AsSequence for PyWeakProxy {
356    fn as_sequence() -> &'static PySequenceMethods {
357        static AS_SEQUENCE: LazyLock<PySequenceMethods> = LazyLock::new(|| PySequenceMethods {
358            length: atomic_func!(|seq, vm| PyWeakProxy::sequence_downcast(seq).len(vm)),
359            contains: atomic_func!(|seq, needle, vm| {
360                PyWeakProxy::sequence_downcast(seq).__contains__(needle, vm)
361            }),
362            ..PySequenceMethods::NOT_IMPLEMENTED
363        });
364        &AS_SEQUENCE
365    }
366}
367
368impl AsMapping for PyWeakProxy {
369    fn as_mapping() -> &'static PyMappingMethods {
370        static AS_MAPPING: PyMappingMethods = PyMappingMethods {
371            length: atomic_func!(|mapping, vm| PyWeakProxy::mapping_downcast(mapping).len(vm)),
372            subscript: atomic_func!(|mapping, needle, vm| {
373                PyWeakProxy::mapping_downcast(mapping).getitem(needle, vm)
374            }),
375            ass_subscript: atomic_func!(|mapping, needle, value, vm| {
376                let zelf = PyWeakProxy::mapping_downcast(mapping);
377                if let Some(value) = value {
378                    zelf.setitem(needle, value, vm)
379                } else {
380                    zelf.delitem(needle, vm)
381                }
382            }),
383        };
384        &AS_MAPPING
385    }
386}
387
388impl Representable for PyWeakProxy {
389    #[inline]
390    fn repr(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyStrRef> {
391        zelf.try_upgrade(vm)?.repr(vm)
392    }
393
394    #[cold]
395    fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
396        unreachable!("use repr instead")
397    }
398}
399
400pub(crate) fn init(context: &'static Context) {
401    PyWeakProxy::extend_class(context, context.types.weakproxy_type);
402    PyWeakCallableProxy::extend_class(context, context.types.weakcallableproxy_type);
403}
404
405impl Hashable for PyWeakProxy {
406    fn hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyHash> {
407        zelf.try_upgrade(vm)?.hash(vm)
408    }
409}