RustPython · OddCoincidence · Mar 10, 2019 · Mar 9, 2019 · Mar 10, 2019
diff --git a/vm/src/obj/mod.rs b/vm/src/obj/mod.rs
@@ -31,4 +31,5 @@ pub mod objstr;
 pub mod objsuper;
 pub mod objtuple;
 pub mod objtype;
+pub mod objweakref;
 pub mod objzip;
diff --git a/vm/src/obj/objweakref.rs b/vm/src/obj/objweakref.rs
@@ -0,0 +1,50 @@
+use crate::function::PyRef;
+use crate::obj::objtype::PyClassRef;
+use crate::pyobject::PyObjectPayload2;
+use crate::pyobject::{PyContext, PyObject, PyObjectRef, PyResult};
+use crate::vm::VirtualMachine;
+
+use std::rc::{Rc, Weak};
+
+#[derive(Debug)]
+pub struct PyWeak {
+    referent: Weak<PyObject>,
+}
+
+impl PyWeak {
+    pub fn downgrade(obj: PyObjectRef) -> PyWeak {
+        PyWeak {
+            referent: Rc::downgrade(&obj),
+        }
+    }
+
+    pub fn upgrade(&self) -> Option<PyObjectRef> {
+        self.referent.upgrade()
+    }
+}
+
+impl PyObjectPayload2 for PyWeak {
+    fn required_type(ctx: &PyContext) -> PyObjectRef {
+        ctx.weakref_type()
+    }
+}
+
+pub type PyWeakRef = PyRef<PyWeak>;
+
+impl PyWeakRef {
+    // TODO callbacks
+    fn create(cls: PyClassRef, referent: PyObjectRef, vm: &mut VirtualMachine) -> PyResult<Self> {
+        Self::new_with_type(vm, PyWeak::downgrade(referent), cls)
+    }
+
+    fn call(self, vm: &mut VirtualMachine) -> PyObjectRef {
+        self.referent.upgrade().unwrap_or_else(|| vm.get_none())
+    }
+}
+
+pub fn init(context: &PyContext) {
+    extend_class!(context, &context.weakref_type, {
+        "__new__" => context.new_rustfunc(PyWeakRef::create),
+        "__call__" => context.new_rustfunc(PyWeakRef::call)
+    });
+}
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use std::fmt;
 use std::iter;
 use std::ops::RangeInclusive;
-use std::rc::{Rc, Weak};
+use std::rc::Rc;
 
 use num_bigint::BigInt;
 use num_bigint::ToBigInt;
@@ -43,6 +43,7 @@ use crate::obj::objstr;
 use crate::obj::objsuper;
 use crate::obj::objtuple::{self, PyTuple};
 use crate::obj::objtype::{self, PyClass};
+use crate::obj::objweakref;
 use crate::obj::objzip;
 use crate::vm::VirtualMachine;
 
@@ -67,9 +68,6 @@ Basically reference counting, but then done by rust.
 /// to the python object by 1.
 pub type PyObjectRef = Rc<PyObject>;
 
-/// Same as PyObjectRef, except for being a weak reference.
-pub type PyObjectWeakRef = Weak<PyObject>;
-
 /// Use this type for function which return a python object or and exception.
 /// Both the python object and the python exception are `PyObjectRef` types
 /// since exceptions are also python objects.
@@ -141,6 +139,7 @@ pub struct PyContext {
     pub readonly_property_type: PyObjectRef,
     pub module_type: PyObjectRef,
     pub bound_method_type: PyObjectRef,
+    pub weakref_type: PyObjectRef,
     pub object: PyObjectRef,
     pub exceptions: exceptions::ExceptionZoo,
 }
@@ -185,6 +184,7 @@ impl PyContext {
         let readonly_property_type =
             create_type("readonly_property", &type_type, &object_type, &dict_type);
         let super_type = create_type("super", &type_type, &object_type, &dict_type);
+        let weakref_type = create_type("ref", &type_type, &object_type, &dict_type);
         let generator_type = create_type("generator", &type_type, &object_type, &dict_type);
         let bound_method_type = create_type("method", &type_type, &object_type, &dict_type);
         let str_type = create_type("str", &type_type, &object_type, &dict_type);
@@ -284,6 +284,7 @@ impl PyContext {
             generator_type,
             module_type,
             bound_method_type,
+            weakref_type,
             type_type,
             exceptions,
         };
@@ -316,6 +317,7 @@ impl PyContext {
         objbool::init(&context);
         objcode::init(&context);
         objframe::init(&context);
+        objweakref::init(&context);
         objnone::init(&context);
         objmodule::init(&context);
         exceptions::init(&context);
@@ -450,6 +452,10 @@ impl PyContext {
         self.bound_method_type.clone()
     }
 
+    pub fn weakref_type(&self) -> PyObjectRef {
+        self.weakref_type.clone()
+    }
+
     pub fn type_type(&self) -> PyObjectRef {
         self.type_type.clone()
     }
@@ -465,6 +471,7 @@ impl PyContext {
     pub fn not_implemented(&self) -> PyObjectRef {
         self.not_implemented.clone()
     }
+
     pub fn object(&self) -> PyObjectRef {
         self.object.clone()
     }
@@ -1481,7 +1488,6 @@ into_py_native_func_tuple!((a, A), (b, B), (c, C), (d, D), (e, E));
 /// of rust data for a particular python object. Determine the python type
 /// by using for example the `.typ()` method on a python object.
 pub enum PyObjectPayload {
-    WeakRef { referent: PyObjectWeakRef },
     AnyRustValue { value: Box<dyn std::any::Any> },
 }
 
@@ -1510,7 +1516,6 @@ impl PyObjectPayload2 for PyIteratorValue {
 impl fmt::Debug for PyObjectPayload {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
-            PyObjectPayload::WeakRef { .. } => write!(f, "weakref"),
             PyObjectPayload::AnyRustValue { value } => value.fmt(f),
         }
     }
@@ -1532,11 +1537,8 @@ impl PyObject {
     }
 
     pub fn payload<T: PyObjectPayload2>(&self) -> Option<&T> {
-        if let PyObjectPayload::AnyRustValue { ref value } = self.payload {
-            value.downcast_ref()
-        } else {
-            None
-        }
+        let PyObjectPayload::AnyRustValue { ref value } = self.payload;
+        value.downcast_ref()
     }
 }
 

diff --git a/vm/src/stdlib/re.rs b/vm/src/stdlib/re.rs
@@ -191,20 +191,19 @@ fn match_end(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 
 /// Retrieve inner rust regex from python object:
 fn get_regex<'a>(obj: &'a PyObjectRef) -> &'a Regex {
-    if let PyObjectPayload::AnyRustValue { ref value } = obj.payload {
-        if let Some(regex) = value.downcast_ref::<Regex>() {
-            return regex;
-        }
+    // TODO: Regex shouldn't be stored in payload directly, create newtype wrapper
+    let PyObjectPayload::AnyRustValue { ref value } = obj.payload;
+    if let Some(regex) = value.downcast_ref::<Regex>() {
+        return regex;
     }
     panic!("Inner error getting regex {:?}", obj);
 }
 
 /// Retrieve inner rust match from python object:
 fn get_match<'a>(obj: &'a PyObjectRef) -> &'a PyMatch {
-    if let PyObjectPayload::AnyRustValue { ref value } = obj.payload {
-        if let Some(value) = value.downcast_ref::<PyMatch>() {
-            return value;
-        }
+    let PyObjectPayload::AnyRustValue { ref value } = obj.payload;
+    if let Some(value) = value.downcast_ref::<PyMatch>() {
+        return value;
     }
     panic!("Inner error getting match {:?}", obj);
 }
diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs
@@ -127,10 +127,9 @@ impl Socket {
 }
 
 fn get_socket<'a>(obj: &'a PyObjectRef) -> impl DerefMut<Target = Socket> + 'a {
-    if let PyObjectPayload::AnyRustValue { ref value } = obj.payload {
-        if let Some(socket) = value.downcast_ref::<RefCell<Socket>>() {
-            return socket.borrow_mut();
-        }
+    let PyObjectPayload::AnyRustValue { ref value } = obj.payload;
+    if let Some(socket) = value.downcast_ref::<RefCell<Socket>>() {
+        return socket.borrow_mut();
     }
     panic!("Inner error getting socket {:?}", obj);
 }

diff --git a/vm/src/stdlib/weakref.rs b/vm/src/stdlib/weakref.rs
@@ -5,51 +5,10 @@
 //! - [rust weak struct](https://doc.rust-lang.org/std/rc/struct.Weak.html)
 //!
 
-use crate::pyobject::{
-    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyObjectWeakRef, PyResult,
-    TypeProtocol,
-};
-use crate::VirtualMachine;
-use std::rc::Rc;
+use super::super::pyobject::{PyContext, PyObjectRef};
 
 pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
-    let py_ref_class = py_class!(ctx, "ref", ctx.object(), {
-        "__new__" => ctx.new_rustfunc(ref_new),
-        "__call__" => ctx.new_rustfunc(ref_call)
-    });
-
     py_module!(ctx, "_weakref", {
-        "ref" => py_ref_class
+        "ref" => ctx.weakref_type()
     })
 }
-
-fn ref_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    // TODO: check first argument for subclass of `ref`.
-    arg_check!(vm, args, required = [(cls, None), (referent, None)]);
-    let referent = Rc::downgrade(referent);
-    Ok(PyObject::new(
-        PyObjectPayload::WeakRef { referent },
-        cls.clone(),
-    ))
-}
-
-/// Dereference the weakref, and check if we still refer something.
-fn ref_call(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    // TODO: check first argument for subclass of `ref`.
-    arg_check!(vm, args, required = [(cls, None)]);
-    let referent = get_value(cls);
-    let py_obj = if let Some(obj) = referent.upgrade() {
-        obj
-    } else {
-        vm.get_none()
-    };
-    Ok(py_obj)
-}
-
-fn get_value(obj: &PyObjectRef) -> PyObjectWeakRef {
-    if let PyObjectPayload::WeakRef { referent } = &obj.payload {
-        referent.clone()
-    } else {
-        panic!("Inner error getting weak ref {:?}", obj);
-    }
-}