Skip to content

Commit d533a29

Browse files
skinnyBatskinnyBat
committed
Add proper weakref type
1 parent 0ec034d commit d533a29

File tree

4 files changed

+66
-52
lines changed

4 files changed

+66
-52
lines changed

‎vm/src/obj/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,4 +30,5 @@ pub mod objstr;
3030
pub mod objsuper;
3131
pub mod objtuple;
3232
pub mod objtype;
33+
pub mod objweakref;
3334
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 & 8 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;
@@ -42,6 +42,7 @@ use crate::obj::objstr;
4242
use crate::obj::objsuper;
4343
use crate::obj::objtuple::{self, PyTuple};
4444
use crate::obj::objtype::{self, PyClass};
45+
use crate::obj::objweakref;
4546
use crate::obj::objzip;
4647
use crate::vm::VirtualMachine;
4748

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

69-
/// Same as PyObjectRef, except for being a weak reference.
70-
pub type PyObjectWeakRef = Weak<PyObject>;
71-
7270
/// Use this type for function which return a python object or and exception.
7371
/// Both the python object and the python exception are `PyObjectRef` types
7472
/// since exceptions are also python objects.
@@ -151,6 +149,7 @@ pub struct PyContext {
151149
pub readonly_property_type: PyObjectRef,
152150
pub module_type: PyObjectRef,
153151
pub bound_method_type: PyObjectRef,
152+
pub weakref_type: PyObjectRef,
154153
pub object: PyObjectRef,
155154
pub exceptions: exceptions::ExceptionZoo,
156155
}
@@ -195,6 +194,7 @@ impl PyContext {
195194
let readonly_property_type =
196195
create_type("readonly_property", &type_type, &object_type, &dict_type);
197196
let super_type = create_type("super", &type_type, &object_type, &dict_type);
197+
let weakref_type = create_type("ref", &type_type, &object_type, &dict_type);
198198
let generator_type = create_type("generator", &type_type, &object_type, &dict_type);
199199
let bound_method_type = create_type("method", &type_type, &object_type, &dict_type);
200200
let str_type = create_type("str", &type_type, &object_type, &dict_type);
@@ -294,6 +294,7 @@ impl PyContext {
294294
generator_type,
295295
module_type,
296296
bound_method_type,
297+
weakref_type,
297298
type_type,
298299
exceptions,
299300
};
@@ -326,6 +327,7 @@ impl PyContext {
326327
objbool::init(&context);
327328
objcode::init(&context);
328329
objframe::init(&context);
330+
objweakref::init(&context);
329331
objnone::init(&context);
330332
objmodule::init(&context);
331333
exceptions::init(&context);
@@ -460,6 +462,10 @@ impl PyContext {
460462
self.bound_method_type.clone()
461463
}
462464

465+
pub fn weakref_type(&self) -> PyObjectRef {
466+
self.weakref_type.clone()
467+
}
468+
463469
pub fn type_type(&self) -> PyObjectRef {
464470
self.type_type.clone()
465471
}
@@ -475,6 +481,7 @@ impl PyContext {
475481
pub fn not_implemented(&self) -> PyObjectRef {
476482
self.not_implemented.clone()
477483
}
484+
478485
pub fn object(&self) -> PyObjectRef {
479486
self.object.clone()
480487
}
@@ -1522,9 +1529,6 @@ pub enum PyObjectPayload {
15221529
name: String,
15231530
scope: ScopeRef,
15241531
},
1525-
WeakRef {
1526-
referent: PyObjectWeakRef,
1527-
},
15281532
RustFunction {
15291533
function: PyNativeFunc,
15301534
},
@@ -1545,7 +1549,6 @@ impl fmt::Debug for PyObjectPayload {
15451549
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
15461550
match self {
15471551
PyObjectPayload::MemoryView { ref obj } => write!(f, "bytes/bytearray {:?}", obj),
1548-
PyObjectPayload::WeakRef { .. } => write!(f, "weakref"),
15491552
PyObjectPayload::Iterator { .. } => write!(f, "iterator"),
15501553
PyObjectPayload::EnumerateIterator { .. } => write!(f, "enumerate"),
15511554
PyObjectPayload::FilterIterator { .. } => write!(f, "filter"),

‎vm/src/stdlib/weakref.rs

Lines changed: 4 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -5,50 +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_mod = ctx.new_module("_weakref", ctx.new_scope(None));
17-
18-
let py_ref_class = ctx.new_class("ref", ctx.object());
19-
ctx.set_attr(&py_ref_class, "__new__", ctx.new_rustfunc(ref_new));
20-
ctx.set_attr(&py_ref_class, "__call__", ctx.new_rustfunc(ref_call));
21-
ctx.set_attr(&py_mod, "ref", py_ref_class);
22-
py_mod
23-
}
24-
25-
fn ref_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
26-
// TODO: check first argument for subclass of `ref`.
27-
arg_check!(vm, args, required = [(cls, None), (referent, None)]);
28-
let referent = Rc::downgrade(referent);
29-
Ok(PyObject::new(
30-
PyObjectPayload::WeakRef { referent },
31-
cls.clone(),
32-
))
33-
}
34-
35-
/// Dereference the weakref, and check if we still refer something.
36-
fn ref_call(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
37-
// TODO: check first argument for subclass of `ref`.
38-
arg_check!(vm, args, required = [(cls, None)]);
39-
let referent = get_value(cls);
40-
let py_obj = if let Some(obj) = referent.upgrade() {
41-
obj
42-
} else {
43-
vm.get_none()
44-
};
45-
Ok(py_obj)
46-
}
47-
48-
fn get_value(obj: &PyObjectRef) -> PyObjectWeakRef {
49-
if let PyObjectPayload::WeakRef { referent } = &obj.payload {
50-
referent.clone()
51-
} else {
52-
panic!("Inner error getting weak ref {:?}", obj);
53-
}
11+
py_module!(ctx, "_weakref", {
12+
"ref" => ctx.weakref_type()
13+
})
5414
}

0 commit comments

Comments
 (0)