Skip to content

Commit 7c5be7d

Browse files
committed
Add proper weakref type
1 parent ded7acc commit 7c5be7d

File tree

4 files changed

+64
-49
lines changed

4 files changed

+64
-49
lines changed

vm/src/obj/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,4 +31,5 @@ pub mod objstr;
3131
pub mod objsuper;
3232
pub mod objtuple;
3333
pub mod objtype;
34+
pub mod objweakref;
3435
pub mod objzip;

vm/src/obj/objweakref.rs

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
use crate::function::PyRef;
2+
use crate::obj::objtype::PyClassRef;
3+
use crate::pyobject::PyObjectPayload2;
4+
use crate::pyobject::{PyContext, PyObject, PyObjectRef, PyResult};
5+
use crate::vm::VirtualMachine;
6+
7+
use std::rc::{Rc, Weak};
8+
9+
#[derive(Debug)]
10+
pub struct PyWeak {
11+
referent: Weak<PyObject>,
12+
}
13+
14+
impl PyWeak {
15+
pub fn downgrade(obj: PyObjectRef) -> PyWeak {
16+
PyWeak {
17+
referent: Rc::downgrade(&obj),
18+
}
19+
}
20+
21+
pub fn upgrade(&self) -> Option<PyObjectRef> {
22+
self.referent.upgrade()
23+
}
24+
}
25+
26+
impl PyObjectPayload2 for PyWeak {
27+
fn required_type(ctx: &PyContext) -> PyObjectRef {
28+
ctx.weakref_type()
29+
}
30+
}
31+
32+
pub type PyWeakRef = PyRef<PyWeak>;
33+
34+
impl PyWeakRef {
35+
// TODO callbacks
36+
fn create(cls: PyClassRef, referent: PyObjectRef, vm: &mut VirtualMachine) -> PyResult<Self> {
37+
Self::new_with_type(vm, PyWeak::downgrade(referent), cls)
38+
}
39+
40+
fn call(self, vm: &mut VirtualMachine) -> PyObjectRef {
41+
self.referent.upgrade().unwrap_or_else(|| vm.get_none())
42+
}
43+
}
44+
45+
pub fn init(context: &PyContext) {
46+
extend_class!(context, &context.weakref_type, {
47+
"__new__" => context.new_rustfunc(PyWeakRef::create),
48+
"__call__" => context.new_rustfunc(PyWeakRef::call)
49+
});
50+
}

vm/src/pyobject.rs

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use std::collections::HashMap;
33
use std::fmt;
44
use std::iter;
55
use std::ops::RangeInclusive;
6-
use std::rc::{Rc, Weak};
6+
use std::rc::Rc;
77

88
use num_bigint::BigInt;
99
use num_bigint::ToBigInt;
@@ -43,6 +43,7 @@ use crate::obj::objstr;
4343
use crate::obj::objsuper;
4444
use crate::obj::objtuple::{self, PyTuple};
4545
use crate::obj::objtype::{self, PyClass};
46+
use crate::obj::objweakref;
4647
use crate::obj::objzip;
4748
use crate::vm::VirtualMachine;
4849

@@ -67,9 +68,6 @@ Basically reference counting, but then done by rust.
6768
/// to the python object by 1.
6869
pub type PyObjectRef = Rc<PyObject>;
6970

70-
/// Same as PyObjectRef, except for being a weak reference.
71-
pub type PyObjectWeakRef = Weak<PyObject>;
72-
7371
/// Use this type for function which return a python object or and exception.
7472
/// Both the python object and the python exception are `PyObjectRef` types
7573
/// since exceptions are also python objects.
@@ -141,6 +139,7 @@ pub struct PyContext {
141139
pub readonly_property_type: PyObjectRef,
142140
pub module_type: PyObjectRef,
143141
pub bound_method_type: PyObjectRef,
142+
pub weakref_type: PyObjectRef,
144143
pub object: PyObjectRef,
145144
pub exceptions: exceptions::ExceptionZoo,
146145
}
@@ -185,6 +184,7 @@ impl PyContext {
185184
let readonly_property_type =
186185
create_type("readonly_property", &type_type, &object_type, &dict_type);
187186
let super_type = create_type("super", &type_type, &object_type, &dict_type);
187+
let weakref_type = create_type("ref", &type_type, &object_type, &dict_type);
188188
let generator_type = create_type("generator", &type_type, &object_type, &dict_type);
189189
let bound_method_type = create_type("method", &type_type, &object_type, &dict_type);
190190
let str_type = create_type("str", &type_type, &object_type, &dict_type);
@@ -284,6 +284,7 @@ impl PyContext {
284284
generator_type,
285285
module_type,
286286
bound_method_type,
287+
weakref_type,
287288
type_type,
288289
exceptions,
289290
};
@@ -316,6 +317,7 @@ impl PyContext {
316317
objbool::init(&context);
317318
objcode::init(&context);
318319
objframe::init(&context);
320+
objweakref::init(&context);
319321
objnone::init(&context);
320322
objmodule::init(&context);
321323
exceptions::init(&context);
@@ -450,6 +452,10 @@ impl PyContext {
450452
self.bound_method_type.clone()
451453
}
452454

455+
pub fn weakref_type(&self) -> PyObjectRef {
456+
self.weakref_type.clone()
457+
}
458+
453459
pub fn type_type(&self) -> PyObjectRef {
454460
self.type_type.clone()
455461
}
@@ -465,6 +471,7 @@ impl PyContext {
465471
pub fn not_implemented(&self) -> PyObjectRef {
466472
self.not_implemented.clone()
467473
}
474+
468475
pub fn object(&self) -> PyObjectRef {
469476
self.object.clone()
470477
}
@@ -1481,7 +1488,6 @@ into_py_native_func_tuple!((a, A), (b, B), (c, C), (d, D), (e, E));
14811488
/// of rust data for a particular python object. Determine the python type
14821489
/// by using for example the `.typ()` method on a python object.
14831490
pub enum PyObjectPayload {
1484-
WeakRef { referent: PyObjectWeakRef },
14851491
AnyRustValue { value: Box<dyn std::any::Any> },
14861492
}
14871493

@@ -1510,7 +1516,6 @@ impl PyObjectPayload2 for PyIteratorValue {
15101516
impl fmt::Debug for PyObjectPayload {
15111517
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
15121518
match self {
1513-
PyObjectPayload::WeakRef { .. } => write!(f, "weakref"),
15141519
PyObjectPayload::AnyRustValue { value } => value.fmt(f),
15151520
}
15161521
}

vm/src/stdlib/weakref.rs

Lines changed: 2 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -5,51 +5,10 @@
55
//! - [rust weak struct](https://doc.rust-lang.org/std/rc/struct.Weak.html)
66
//!
77
8-
use crate::pyobject::{
9-
PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyObjectWeakRef, PyResult,
10-
TypeProtocol,
11-
};
12-
use crate::VirtualMachine;
13-
use std::rc::Rc;
8+
use super::super::pyobject::{PyContext, PyObjectRef};
149

1510
pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
16-
let py_ref_class = py_class!(ctx, "ref", ctx.object(), {
17-
"__new__" => ctx.new_rustfunc(ref_new),
18-
"__call__" => ctx.new_rustfunc(ref_call)
19-
});
20-
2111
py_module!(ctx, "_weakref", {
22-
"ref" => py_ref_class
12+
"ref" => ctx.weakref_type()
2313
})
2414
}
25-
26-
fn ref_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
27-
// TODO: check first argument for subclass of `ref`.
28-
arg_check!(vm, args, required = [(cls, None), (referent, None)]);
29-
let referent = Rc::downgrade(referent);
30-
Ok(PyObject::new(
31-
PyObjectPayload::WeakRef { referent },
32-
cls.clone(),
33-
))
34-
}
35-
36-
/// Dereference the weakref, and check if we still refer something.
37-
fn ref_call(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
38-
// TODO: check first argument for subclass of `ref`.
39-
arg_check!(vm, args, required = [(cls, None)]);
40-
let referent = get_value(cls);
41-
let py_obj = if let Some(obj) = referent.upgrade() {
42-
obj
43-
} else {
44-
vm.get_none()
45-
};
46-
Ok(py_obj)
47-
}
48-
49-
fn get_value(obj: &PyObjectRef) -> PyObjectWeakRef {
50-
if let PyObjectPayload::WeakRef { referent } = &obj.payload {
51-
referent.clone()
52-
} else {
53-
panic!("Inner error getting weak ref {:?}", obj);
54-
}
55-
}

0 commit comments

Comments
 (0)