Commit 387cdc5d authored by Joseph Weston's avatar Joseph Weston Committed by Christoph Groth
Browse files

implement pickling

parent 25520bf3
......@@ -1486,6 +1486,109 @@ Array<T> *Array<T>::make(int ndim, const size_t *shape, size_t *sizep)
return result;
}
static bool little_endian()
{
union {
uint32_t i;
char c[4];
} bint = {0x04030201};
return bint.c[0] == 1;
}
template <typename T>
PyObject *tinyarray_reduce(PyObject *py_self)
{
PyObject *py_ret, *py_state, *py_constructor, *py_mod;
assert(Array<T>::check_exact(self_)); Array<T> *self = (Array<T>*)py_self;
py_ret = PyTuple_New(3);
if (!py_ret) {
return 0;
}
py_mod = PyImport_ImportModule("tinyarray");
if (!py_mod) {
Py_DECREF(py_ret);
return 0;
}
py_constructor = PyObject_GetAttrString(py_mod, "_empty");
Py_DECREF(py_mod);
int ndim;
size_t *shape;
T *data = self->data();
self->ndim_shape(&ndim, &shape);
size_t size_in_bytes = calc_size(ndim, shape) * sizeof(T);
PyObject *py_data = PyString_FromStringAndSize((char*)data, size_in_bytes);
PyObject *py_dtype = get_dtype_py(py_self, NULL);
PyObject *py_shape = PyTuple_New(ndim);
for (int i=0; i < ndim; ++i)
PyTuple_SET_ITEM(py_shape, i, PyInt_FromSize_t(shape[i]));
PyTuple_SET_ITEM(py_ret, 0, py_constructor);
PyTuple_SET_ITEM(py_ret, 1, Py_BuildValue("(OO)", py_shape, py_dtype));
// fill in state to be passed to __setstate__
py_state = PyTuple_New(3);
if (!py_state) {
Py_DECREF(py_ret);
return 0;
}
PyTuple_SET_ITEM(py_state, 0, py_data);
PyTuple_SET_ITEM(py_state, 1, PyInt_FromLong(little_endian()));
PyTuple_SET_ITEM(py_state, 2, PyInt_FromLong(sizeof(int)));
PyTuple_SET_ITEM(py_ret, 2, py_state);
return py_ret;
}
template <typename T>
PyObject *tinyarray_reduce_ex(PyObject *py_self, PyObject*)
{
return tinyarray_reduce<T>(py_self);
}
template <typename T>
PyObject *tinyarray_setstate(PyObject *py_self, PyObject *py_args)
{
assert(Array<T>::check_exact(self_)); Array<T> *self = (Array<T>*)py_self;
T* data = 0;
size_t size_in_bytes;
Dtype dtype = get_dtype(py_self);
int pickle_little_endian, pickle_byte_size;
PyArg_ParseTuple(py_args, "(s#ii)",
(char*)&data, &size_in_bytes,
&pickle_little_endian,
&pickle_byte_size);
// Check for endianness and integer size between pickling and unpickling
// environment.
if (dtype == Dtype::LONG && little_endian() != pickle_little_endian) {
PyErr_SetString(PyExc_RuntimeError,
"Endianness of pickled TinyArray and "
"local machine are not identical.");
return 0;
}
if (dtype == Dtype::LONG && pickle_byte_size != sizeof(int)) {
PyErr_SetString(PyExc_RuntimeError,
"integer size of pickled TinyArray is not the same "
"as the integer size of local machine.");
return 0;
}
size_t size = size_in_bytes / sizeof(T);
T *self_data = self->data();
for(size_t i = 0; i < size; ++i) self_data[i] = data[i];
Py_INCREF(Py_None);
return Py_None;
}
// **************** Type object ****************
template <>
......@@ -1523,6 +1626,9 @@ template <typename T>
PyMethodDef Array<T>::methods[] = {
{"transpose", (PyCFunction)transpose<T>, METH_NOARGS},
{"conjugate", (PyCFunction)apply_unary_ufunc<Conjugate<T> >, METH_NOARGS},
{"__reduce__", (PyCFunction)tinyarray_reduce<T>, METH_NOARGS},
{"__reduce_ex__", (PyCFunction)tinyarray_reduce_ex<T>, METH_VARARGS},
{"__setstate__", (PyCFunction)tinyarray_setstate<T>, METH_VARARGS},
{0, 0} // Sentinel
};
......@@ -1579,3 +1685,15 @@ template class Array<Complex>;
template PyObject *transpose<long>(PyObject*);
template PyObject *transpose<double>(PyObject*);
template PyObject *transpose<Complex>(PyObject*);
template PyObject *tinyarray_reduce<long>(PyObject*);
template PyObject *tinyarray_reduce<double>(PyObject*);
template PyObject *tinyarray_reduce<Complex>(PyObject*);
template PyObject *tinyarray_reduce_ex<long>(PyObject*, PyObject*);
template PyObject *tinyarray_reduce_ex<double>(PyObject*, PyObject*);
template PyObject *tinyarray_reduce_ex<Complex>(PyObject*, PyObject*);
template PyObject *tinyarray_setstate<long>(PyObject*, PyObject*);
template PyObject *tinyarray_setstate<double>(PyObject*, PyObject*);
template PyObject *tinyarray_setstate<Complex>(PyObject*, PyObject*);
......@@ -23,6 +23,17 @@ int dtype_converter(const PyObject *ob, Dtype *dtype)
return 1;
}
template <typename T>
PyObject *unfilled(int ndim, const size_t *shape)
{
size_t size ;
Array<T> *result = Array<T>::make(ndim, shape, &size);
return result ? (PyObject*)result : 0 ;
}
PyObject *(*unfilled_dtable[])(int, const size_t*) =
DTYPE_DISPATCH(unfilled);
template <typename T>
PyObject *filled(int ndim, const size_t *shape, int value)
{
......@@ -54,6 +65,23 @@ PyObject *filled_pyargs(PyObject *args, int value)
return filled_dtable[int(dtype)](ndim, shape, value);
}
PyObject *empty(PyObject *, PyObject *args)
{
PyObject *pyshape;
Dtype dtype = default_dtype;
if (!PyArg_ParseTuple(args, "O|O&", &pyshape, dtype_converter, &dtype))
return 0;
unsigned long shape_as_ulong[max_ndim];
int ndim = load_index_seq_as_ulong(pyshape, shape_as_ulong, max_ndim,
"Negative dimensions are not allowed.");
if (ndim == -1) return 0;
size_t shape[max_ndim];
for (int d = 0; d < ndim; ++d) shape[d] = shape_as_ulong[d];
return unfilled_dtable[int(dtype)](ndim, shape);
}
PyObject *zeros(PyObject *, PyObject *args)
{
return filled_pyargs(args, 0);
......@@ -198,6 +226,7 @@ PyObject *unary_ufunc_round(PyObject *, PyObject *args)
// template <typename T> using Round_ceil = Round<Ceil, T>;
PyMethodDef functions[] = {
{"_empty", empty, METH_VARARGS},
{"zeros", zeros, METH_VARARGS},
{"ones", ones, METH_VARARGS},
{"identity", identity, METH_VARARGS},
......
......@@ -315,3 +315,24 @@ def test_unary_ufuncs():
-1.7, -1.5, -0.5, -0.3, -0.0, 0.0, 0.3, 0.5, 1.5, 1.7,
2.0, 2.5, 3.51, 4.51, 987654321.5, 987654322.5]:
assert_equal(ta_func(x), np_func(x))
def test_pickle():
try:
import cPickle as pickle
except ImportError:
import pickle
import numpy.random as nr
tests = [ta.array((100 * nr.rand(10)).astype(int)), # 1D
ta.array(nr.rand(10)),
ta.array(nr.rand(10) + 1j * nr.rand(10)),
ta.array((100 * nr.rand(10, 5)).astype(int)), # 2D
ta.array(nr.rand(10, 5)),
ta.array(nr.rand(10, 5) + 1j * nr.rand(10, 5)),
ta.array((100 * nr.rand(10, 5, 3)).astype(int)), # 3D
ta.array(nr.rand(10, 5, 3)),
ta.array(nr.rand(10, 5, 3) + 1j * nr.rand(10, 5, 3))]
results = [pickle.loads(pickle.dumps(t)) for t in tests]
[assert_equal(t, r) for t, r in zip(tests, results)]
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