impl __sizeof__ for bytearray, int, list, str, dict and set (RustPython#1426)

Author: RobertBerglund

vm/src/dictdatatype.rs

@@ -9,6 +9,7 @@ use num_bigint::ToBigInt;
 /// And: http://code.activestate.com/recipes/578375/
 use std::collections::{hash_map::DefaultHasher, HashMap};
 use std::hash::{Hash, Hasher};
+use std::mem::size_of;
 
 /// hash value of an object returned by __hash__
 type HashValue = pyhash::PyHash;
@@ -269,6 +270,10 @@ impl<T: Clone> Dict<T> {
             None => None,
         }
     }
+
+    pub fn sizeof(&self) -> usize {
+        size_of::<Self>() + self.size * size_of::<DictEntry<T>>()
+    }
 }
 
 enum LookupResult {

vm/src/obj/objbytearray.rs

@@ -21,6 +21,7 @@ use crate::pyobject::{
     TryFromObject,
 };
 use crate::vm::VirtualMachine;
+use std::mem::size_of;
 
 /// "bytearray(iterable_of_ints) -> bytearray\n\
 ///  bytearray(string, encoding[, errors]) -> bytearray\n\
@@ -109,6 +110,11 @@ impl PyByteArrayRef {
         self.inner.borrow().len()
     }
 
+    #[pymethod(name = "__sizeof__")]
+    fn sizeof(self, _vm: &VirtualMachine) -> usize {
+        size_of::<Self>() + self.inner.borrow().len() * size_of::<u8>()
+    }
+
     #[pymethod(name = "__eq__")]
     fn eq(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         self.inner.borrow().eq(other, vm)

vm/src/obj/objdict.rs

@@ -12,6 +12,8 @@ use crate::pyobject::{
 };
 use crate::vm::{ReprGuard, VirtualMachine};
 
+use std::mem::size_of;
+
 pub type DictContentType = dictdatatype::Dict;
 
 #[derive(Default)]
@@ -154,6 +156,10 @@ impl PyDictRef {
         self.entries.borrow().len()
     }
 
+    fn sizeof(self, _vm: &VirtualMachine) -> usize {
+        size_of::<Self>() + self.entries.borrow().sizeof()
+    }
+
     fn repr(self, vm: &VirtualMachine) -> PyResult<String> {
         let s = if let Some(_guard) = ReprGuard::enter(self.as_object()) {
             let mut str_parts = vec![];
@@ -576,6 +582,7 @@ pub fn init(context: &PyContext) {
     extend_class!(context, &context.types.dict_type, {
         "__bool__" => context.new_rustfunc(PyDictRef::bool),
         "__len__" => context.new_rustfunc(PyDictRef::len),
+        "__sizeof__" => context.new_rustfunc(PyDictRef::sizeof),
         "__contains__" => context.new_rustfunc(PyDictRef::contains),
         "__delitem__" => context.new_rustfunc(PyDictRef::inner_delitem),
         "__eq__" => context.new_rustfunc(PyDictRef::eq),

vm/src/obj/objint.rs

@@ -1,4 +1,5 @@
 use std::fmt;
+use std::mem::size_of;
 use std::str;
 
 use num_bigint::{BigInt, Sign};
@@ -573,6 +574,11 @@ impl PyInt {
         !self.value.is_zero()
     }
 
+    #[pymethod(name = "__sizeof__")]
+    fn sizeof(&self, _vm: &VirtualMachine) -> usize {
+        size_of::<Self>()
+    }
+
     #[pymethod]
     fn bit_length(&self, _vm: &VirtualMachine) -> usize {
         self.value.bits()

vm/src/obj/objlist.rs

@@ -1,5 +1,6 @@
 use std::cell::{Cell, RefCell};
 use std::fmt;
+use std::mem::size_of;
 use std::ops::Range;
 
 use num_bigint::{BigInt, ToBigInt};
@@ -191,6 +192,10 @@ impl PyListRef {
         self.elements.borrow().len()
     }
 
+    fn sizeof(self, _vm: &VirtualMachine) -> usize {
+        size_of::<Self>() + self.elements.borrow().capacity() * size_of::<PyObjectRef>()
+    }
+
     fn reverse(self, _vm: &VirtualMachine) {
         self.elements.borrow_mut().reverse();
     }
@@ -879,6 +884,7 @@ pub fn init(context: &PyContext) {
                     The argument must be an iterable if specified.";
 
     extend_class!(context, list_type, {
+        "__sizeof__" => context.new_rustfunc(PyListRef::sizeof),
         "__add__" => context.new_rustfunc(PyListRef::add),
         "__iadd__" => context.new_rustfunc(PyListRef::iadd),
         "__bool__" => context.new_rustfunc(PyListRef::bool),

vm/src/obj/objset.rs

@@ -91,6 +91,10 @@ impl PySetInner {
         self.content.len()
     }
 
+    fn sizeof(&self) -> usize {
+        self.content.sizeof()
+    }
+
     fn copy(&self) -> PySetInner {
         PySetInner {
             content: self.content.clone(),
@@ -343,6 +347,11 @@ impl PySet {
         self.inner.borrow().len()
     }
 
+    #[pymethod(name = "__sizeof__")]
+    fn sizeof(&self, _vm: &VirtualMachine) -> usize {
+        std::mem::size_of::<Self>() + self.inner.borrow().sizeof()
+    }
+
     #[pymethod]
     fn copy(&self, _vm: &VirtualMachine) -> Self {
         Self {
@@ -594,6 +603,11 @@ impl PyFrozenSet {
         self.inner.len()
     }
 
+    #[pymethod(name = "__sizeof__")]
+    fn sizeof(&self, _vm: &VirtualMachine) -> usize {
+        std::mem::size_of::<Self>() + self.inner.sizeof()
+    }
+
     #[pymethod]
     fn copy(&self, _vm: &VirtualMachine) -> Self {
         Self {

vm/src/obj/objstr.rs

@@ -4,6 +4,7 @@ extern crate unicode_xid;
 use std::cell::Cell;
 use std::char;
 use std::fmt;
+use std::mem::size_of;
 use std::ops::Range;
 use std::str::FromStr;
 use std::string::ToString;
@@ -285,6 +286,11 @@ impl PyString {
         self.value.chars().count()
     }
 
+    #[pymethod(name = "__sizeof__")]
+    fn sizeof(&self, _vm: &VirtualMachine) -> usize {
+        size_of::<Self>() + self.value.capacity() * size_of::<u8>()
+    }
+
     #[pymethod(name = "__mul__")]
     fn mul(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult<String> {
         if !objtype::isinstance(&val, &vm.ctx.int_type()) {