Commit fc33353c authored by Joseph Weston's avatar Joseph Weston
Browse files

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.
parent a982cddf
Pipeline #361 passed with stage
......@@ -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*);
......@@ -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
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment