implement __sizeof__ for tinyarray

A new method `object_size` is added to the Array template class, which
calculates the total size of the Python object in bytes.  The function
`size_of` then wraps this method and provides the necessary interface
for a Python method object.

src/array.cc

@@ -1562,6 +1562,13 @@ PyObject *conjugate(PyObject *in_, PyObject *)
   return apply_unary_ufunc<Conjugate<T> >(in_);
 }
 
+template <typename T>
+PyObject *size_of(PyObject *in_, PyObject *)
+{
+    assert(Array<T>::check_exact(in_));
+    return PyInt_FromSsize_t(((const Array<T>*) in_)->object_size());
+}
+
 template <typename T>
 Array<T> *Array<T>::make(int ndim, size_t size)
 {
@@ -1605,6 +1612,7 @@ Array<T> *Array<T>::make(int ndim, const size_t *shape, size_t *sizep)
     return result;
 }
 
+
 template <typename T>
 PyObject *reduce(PyObject *self_, PyObject*)
 {
@@ -1675,6 +1683,7 @@ template <typename T>
 PyMethodDef Array<T>::methods[] = {
     {"transpose", (PyCFunction)transpose<T>, METH_NOARGS},
     {"conjugate", (PyCFunction)conjugate<T>, METH_NOARGS},
+    {"__sizeof__", (PyCFunction)size_of<T>, METH_NOARGS},
     {"__reduce__", (PyCFunction)reduce<T>, METH_NOARGS},
     {0, 0}                      // Sentinel
 };
@@ -1743,6 +1752,10 @@ template PyObject *conjugate<long>(PyObject*, PyObject*);
 template PyObject *conjugate<double>(PyObject*, PyObject*);
 template PyObject *conjugate<Complex>(PyObject*, PyObject*);
 
+template PyObject *size_of<long>(PyObject*, PyObject*);
+template PyObject *size_of<double>(PyObject*, PyObject*);
+template PyObject *size_of<Complex>(PyObject*, PyObject*);
+
 template PyObject *reduce<long>(PyObject*, PyObject*);
 template PyObject *reduce<double>(PyObject*, PyObject*);
 template PyObject *reduce<Complex>(PyObject*, PyObject*);

src/array.hh

@@ -105,6 +105,8 @@ public:
         }
     }
 
+    ssize_t object_size() const;
+
     static bool check_exact(PyObject *candidate) {
         return (Py_TYPE(candidate) == &pytype);
     }
@@ -157,4 +159,23 @@ int coerce_to_arrays(PyObject **a, PyObject **b, Dtype *coerced_dtype);
 
 template <typename T> PyObject *transpose(PyObject *in, PyObject *dummy);
 
+template <typename T>
+ssize_t Array<T>::object_size() const
+{
+    int ndim;
+    size_t *shape;
+    ndim_shape(&ndim, &shape);
+
+    // this does the same calculation as in Array<T>::make
+    Py_ssize_t size = calc_size(ndim, shape);
+    if (ndim > 1)
+        // if array is > 1D then the shape is stored in
+        // the start of the data buffer
+        size += (ndim * sizeof(size_t) + sizeof(T) - 1) / sizeof(T);
+    size = size * sizeof(T);  // element count -> byte count
+    size += Array<T>::pytype.tp_basicsize;  // add Python object overhead
+
+    return size;
+}
+
 #endif // !ARRAY_HH

test_tinyarray.py

@@ -8,6 +8,7 @@
 # https://gitlab.kwant-project.org/kwant/tinyarray.
 
 import operator, warnings
+import platform
 import tinyarray as ta
 from nose.tools import assert_raises
 import numpy as np
@@ -15,8 +16,23 @@ from numpy.testing import assert_equal, assert_almost_equal
 import sys
 import random
 
+def machine_wordsize():
+    bits, _ = platform.architecture()
+    if bits == '32bit':
+        return 4
+    elif bits == '64bit':
+        return 8
+    else:
+        raise RuntimeError('unknown architecture', bits)
+
 dtypes = [int, float, complex]
 
+dtype_size = {
+    int: machine_wordsize(),
+    float: 8,
+    complex: 16
+}
+
 some_shapes = [(), 0, 1, 2, 3,
                (0, 0), (1, 0), (0, 1), (2, 2), (17, 17),
                (0, 0, 0), (1, 1, 1), (2, 2, 1), (2, 0, 3)]
@@ -375,6 +391,27 @@ def test_other_scalar_types():
         assert_equal(ta.matrix(a), np.matrix(a))
 
 
+def test_sizeof():
+    obj = object()
+    word_size = machine_wordsize()
+    for shape in some_shapes:
+        for dtype in dtypes:
+            a = ta.zeros(shape, dtype)
+            sizeof = a.__sizeof__()
+            # basic buffer size
+            n_elements = a.size
+            # if the array is > 1D then the shape is stored
+            # at the start of the buffer
+            if len(a.shape) > 1:
+                n_elements += (a.ndim * machine_wordsize() +
+                               dtype_size[dtype] - 1) // dtype_size[dtype]
+            buffer_size = n_elements * dtype_size[dtype]
+            # basic Python object has 3 pointer-sized members
+            sizeof_should_be = buffer_size + 3 * machine_wordsize()
+            print(dtype, shape)
+            assert_equal(sizeof, sizeof_should_be)
+
+
 def test_pickle():
     import pickle