Move weak_refs attribute into Instance payload

RustPython · BenLewis-Seequent · Feb 13, 2019 · Feb 16, 2019 · Feb 17, 2019 · Feb 18, 2019
commit 350bc30311ef76fcbcfaf34a3521edd5c71eaa34
diff --git a/tests/snippets/weakrefs.py b/tests/snippets/weakrefs.py
@@ -1,6 +1,7 @@
 import sys
 from _weakref import ref
 
+from testutils import assert_raises
 
 data_holder = {}
 
@@ -58,3 +59,4 @@ def never_callback(_weak_ref):
 
     assert str(value) == "param: 5"
 
+assert_raises(TypeError, lambda: ref(1), "can't create weak reference to an int")
diff --git a/vm/src/obj/objtype.rs b/vm/src/obj/objtype.rs
@@ -232,7 +232,7 @@ pub fn get_attributes(obj: &PyObjectRef) -> PyAttributes {
     }
 
     // Get instance attributes:
-    if let PyObjectPayload::Instance { dict } = &obj.borrow().payload {
+    if let PyObjectPayload::Instance { dict, .. } = &obj.borrow().payload {
         for (name, value) in dict.borrow().iter() {
             attributes.insert(name.to_string(), value.clone());
         }

diff --git a/vm/src/obj/objweakref.rs b/vm/src/obj/objweakref.rs
@@ -21,8 +21,15 @@ fn ref_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         },
         cls.clone(),
     );
-    referent.borrow_mut().add_weakref(&weakref);
-    Ok(weakref)
+    if referent.borrow_mut().add_weakref(&weakref) {
+        Ok(weakref)
+    } else {
+        let referent_repr = vm.to_pystr(&referent.typ())?;
+        Err(vm.new_type_error(format!(
+            "cannot create weak reference to '{}' object",
+            referent_repr
+        )))
+    }
 }
 
 /// Dereference the weakref, and check if we still refer something.

diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -92,12 +92,13 @@ impl PyObjectRef {
         while let Some(obj) = DELETE_QUEUE.with(|q| q.borrow_mut().pop_front()) {
             // check is still needs to be deleted
             if PyObjectRef::strong_count(&obj) == 1 {
-                let weak_refs: Vec<PyObjectRef> = obj
-                    .borrow()
-                    .weak_refs
-                    .iter()
-                    .flat_map(|x| x.upgrade())
-                    .collect();
+                let weak_refs = match obj.borrow().payload {
+                    PyObjectPayload::Instance { ref weak_refs, .. } => weak_refs
+                        .iter()
+                        .flat_map(|x| x.upgrade())
+                        .collect::<Vec<_>>(),
+                    _ => panic!("Non-instance object can't have weakrefs"),
+                };
 
                 for weak_ref in weak_refs.iter() {
                     objweakref::clear_weak_ref(weak_ref);
@@ -137,15 +138,22 @@ impl Deref for PyObjectRef {
 
 impl Drop for PyObjectRef {
     fn drop(&mut self) {
-        if PyObjectRef::strong_count(self) == 1 && !self.borrow().weak_refs.is_empty() {
-            // The PyObject is going to be dropped, delete it later when we have access to vm.
-            // This will error when the tls is destroyed, thus meaning objects won't actually
-            // be properly destroyed.
-            if DELETE_QUEUE
-                .try_with(|q| q.borrow_mut().push_back(self.clone()))
-                .is_err()
-            {
-                warn!("Object hasn't been cleaned up {:?}.", self);
+        if PyObjectRef::strong_count(self) == 1 {
+            let has_weakrefs = match self.borrow().payload {
+                PyObjectPayload::Instance { ref weak_refs, .. } => !weak_refs.is_empty(),
+                _ => false,
+            };
+
+            if has_weakrefs {
+                // The PyObject is going to be dropped, delete it later when we have access to vm.
+                // This will error when the tls is destroyed, thus meaning objects won't actually
+                // be properly destroyed.
+                if DELETE_QUEUE
+                    .try_with(|q| q.borrow_mut().push_back(self.clone()))
+                    .is_err()
+                {
+                    warn!("Object hasn't been cleaned up {:?}.", self);
+                }
             }
         }
     }
@@ -738,6 +746,7 @@ impl PyContext {
         PyObject::new(
             PyObjectPayload::Instance {
                 dict: RefCell::new(dict),
+                weak_refs: Vec::new(),
             },
             class,
         )
@@ -764,7 +773,8 @@ impl PyContext {
     pub fn set_attr(&self, obj: &PyObjectRef, attr_name: &str, value: PyObjectRef) {
         match obj.borrow().payload {
             PyObjectPayload::Module { ref dict, .. } => self.set_attr(dict, attr_name, value),
-            PyObjectPayload::Instance { ref dict } | PyObjectPayload::Class { ref dict, .. } => {
+            PyObjectPayload::Instance { ref dict, .. }
+            | PyObjectPayload::Class { ref dict, .. } => {
                 dict.borrow_mut().insert(attr_name.to_string(), value);
             }
             PyObjectPayload::Scope { ref scope } => {
@@ -801,7 +811,6 @@ impl PyContext {
 pub struct PyObject {
     pub payload: PyObjectPayload,
     pub typ: Option<PyObjectRef>,
-    weak_refs: Vec<PyObjectWeakRef>,
     // pub dict: HashMap<String, PyObjectRef>, // __dict__ member
 }
 
@@ -904,7 +913,7 @@ impl AttributeProtocol for PyObjectRef {
                 }
                 None
             }
-            PyObjectPayload::Instance { ref dict } => dict.borrow().get(attr_name).cloned(),
+            PyObjectPayload::Instance { ref dict, .. } => dict.borrow().get(attr_name).cloned(),
             _ => None,
         }
     }
@@ -916,7 +925,7 @@ impl AttributeProtocol for PyObjectRef {
             PyObjectPayload::Class { ref mro, .. } => {
                 class_has_item(self, attr_name) || mro.iter().any(|d| class_has_item(d, attr_name))
             }
-            PyObjectPayload::Instance { ref dict } => dict.borrow().contains_key(attr_name),
+            PyObjectPayload::Instance { ref dict, .. } => dict.borrow().contains_key(attr_name),
             _ => false,
         }
     }
@@ -1124,6 +1133,7 @@ pub enum PyObjectPayload {
     },
     Instance {
         dict: RefCell<PyAttributes>,
+        weak_refs: Vec<PyObjectWeakRef>,
     },
     RustFunction {
         function: Box<Fn(&mut VirtualMachine, PyFuncArgs) -> PyResult>,
@@ -1178,7 +1188,6 @@ impl PyObject {
             payload,
             typ: Some(typ),
             // dict: HashMap::new(),  // dict,
-            weak_refs: Vec::new(),
         }
         .into_ref()
     }
@@ -1188,7 +1197,6 @@ impl PyObject {
             payload,
             typ: None,
             // dict: HashMap::new(),  // dict,
-            weak_refs: Vec::new(),
         }
         .into_ref()
     }
@@ -1200,8 +1208,16 @@ impl PyObject {
         }
     }
 
-    pub fn add_weakref(&mut self, weakref: &PyObjectRef) {
-        self.weak_refs.push(PyObjectRef::downgrade(weakref))
+    pub fn add_weakref(&mut self, weakref: &PyObjectRef) -> bool {
+        match self.payload {
+            PyObjectPayload::Instance {
+                ref mut weak_refs, ..
+            } => {
+                weak_refs.push(PyObjectRef::downgrade(weakref));
+                true
+            }
+            _ => false,
+        }
     }
 }