Merge pull request #5001 from juliantaylor/sizeof

support __sizeof__()

4 files changed, 96 insertions, 0 deletions

diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
index a791c2c22..bfd3bc3c1 100644
--- a/numpy/core/src/multiarray/methods.c
+++ b/numpy/core/src/multiarray/methods.c
@@ -1901,6 +1901,19 @@ array_dumps(PyArrayObject *self, PyObject *args)
 
 
 static PyObject *
+array_sizeof(PyArrayObject *self)
+{
+    /* object + dimension and strides */
+    Py_ssize_t nbytes = NPY_SIZEOF_PYARRAYOBJECT +
+        PyArray_NDIM(self) * sizeof(npy_intp) * 2;
+    if (PyArray_CHKFLAGS(self, NPY_ARRAY_OWNDATA)) {
+        nbytes += PyArray_NBYTES(self);
+    }
+    return PyLong_FromSsize_t(nbytes);
+}
+
+
+static PyObject *
 array_transpose(PyArrayObject *self, PyObject *args)
 {
     PyObject *shape = Py_None;
@@ -2308,6 +2321,11 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = {
         (PyCFunction)array_wraparray,
         METH_VARARGS, NULL},
 
+    /* for the sys module */
+    {"__sizeof__",
+        (PyCFunction) array_sizeof,
+        METH_NOARGS, NULL},
+
     /* for the copy module */
     {"__copy__",
         (PyCFunction)array_copy_keeporder,
diff --git a/numpy/core/src/multiarray/scalartypes.c.src b/numpy/core/src/multiarray/scalartypes.c.src
index 4fa634098..54f57bcf4 100644
--- a/numpy/core/src/multiarray/scalartypes.c.src
+++ b/numpy/core/src/multiarray/scalartypes.c.src
@@ -1193,6 +1193,20 @@ gentype_size_get(PyObject *NPY_UNUSED(self))
     return PyInt_FromLong(1);
 }
 
+static PyObject *
+gentype_sizeof(PyObject *self)
+{
+    Py_ssize_t nbytes;
+    PyObject * isz = gentype_itemsize_get(self);
+    if (isz == NULL) {
+        return NULL;
+    }
+    nbytes = PyLong_AsLong(isz) + Py_TYPE(self)->tp_basicsize +
+        Py_SIZE(self) * Py_TYPE(self)->tp_itemsize;
+    Py_DECREF(isz);
+    return PyLong_FromSsize_t(nbytes);
+}
+
 #if PY_VERSION_HEX >= 0x03000000
 NPY_NO_EXPORT void
 gentype_struct_free(PyObject *ptr)
@@ -1921,6 +1935,11 @@ static PyMethodDef gentype_methods[] = {
         (PyCFunction)gentype_wraparray,
         METH_VARARGS, doc_sc_wraparray},
 
+    /* for the sys module */
+    {"__sizeof__",
+        (PyCFunction)gentype_sizeof,
+        METH_NOARGS, NULL},
+
     /* for the copy module */
     {"__copy__",
         (PyCFunction)gentype_copy,
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 70398ee84..ade95fb7b 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -4566,5 +4566,51 @@ class TestWhere(TestCase):
         assert_equal(np.where(False, b, a), "abcd")
 
 
+class TestSizeOf(TestCase):
+
+    def test_empty_array(self):
+        x = np.array([])
+        assert_(sys.getsizeof(x) > 0)
+
+    def check_array(self, dtype):
+        elem_size = dtype(0).itemsize
+
+        for length in [10, 50, 100, 500]:
+            x = np.arange(length, dtype=dtype)
+            assert_(sys.getsizeof(x) > length * elem_size)
+
+    def test_array_int32(self):
+        self.check_array(np.int32)
+
+    def test_array_int64(self):
+        self.check_array(np.int64)
+
+    def test_array_float32(self):
+        self.check_array(np.float32)
+
+    def test_array_float64(self):
+        self.check_array(np.float64)
+
+    def test_view(self):
+        d = np.ones(100)
+        assert_(sys.getsizeof(d[...]) < sys.getsizeof(d))
+
+    def test_reshape(self):
+        d = np.ones(100)
+        assert_(sys.getsizeof(d) < sys.getsizeof(d.reshape(100, 1, 1).copy()))
+
+    def test_resize(self):
+        d = np.ones(100)
+        old = sys.getsizeof(d)
+        d.resize(50)
+        assert_(old > sys.getsizeof(d))
+        d.resize(150)
+        assert_(old < sys.getsizeof(d))
+
+    def test_error(self):
+        d = np.ones(100)
+        assert_raises(TypeError, d.__sizeof__, "a")
+
+
 if __name__ == "__main__":
     run_module_suite()
diff --git a/numpy/core/tests/test_scalarmath.py b/numpy/core/tests/test_scalarmath.py
index afdc06c03..3ba3beff9 100644
--- a/numpy/core/tests/test_scalarmath.py
+++ b/numpy/core/tests/test_scalarmath.py
@@ -271,5 +271,18 @@ class TestRepr(object):
         for t in [np.float32, np.float64]:
             yield self._test_type_repr, t
 
+
+class TestSizeOf(TestCase):
+
+    def test_equal_nbytes(self):
+        for type in types:
+            x = type(0)
+            assert_(sys.getsizeof(x) > x.nbytes)
+
+    def test_error(self):
+        d = np.float32()
+        assert_raises(TypeError, d.__sizeof__, "a")
+
+
 if __name__ == "__main__":
     run_module_suite()