summaryrefslogtreecommitdiff
path: root/numpy/core/src/arrayobject.c
diff options
context:
space:
mode:
Diffstat (limited to 'numpy/core/src/arrayobject.c')
-rw-r--r--numpy/core/src/arrayobject.c69
1 files changed, 69 insertions, 0 deletions
diff --git a/numpy/core/src/arrayobject.c b/numpy/core/src/arrayobject.c
index ad3cbc4f2..f944a499d 100644
--- a/numpy/core/src/arrayobject.c
+++ b/numpy/core/src/arrayobject.c
@@ -10790,6 +10790,75 @@ static PyTypeObject PyArrayMapIter_Type = {
/** END of Subscript Iterator **/
+/*
+ NUMPY_API
+ Get MultiIterator from array of Python objects and any additional
+
+ PyObject **mps -- array of PyObjects
+ int n - number of PyObjects in the array
+ int nadd - number of additional arrays to include in the
+ iterator.
+
+ Returns a multi-iterator object.
+ */
+static PyObject *
+PyArray_MultiIterFromObjects(PyObject **mps, int n, int nadd, ...)
+{
+ va_list va;
+ PyArrayMultiIterObject *multi;
+ PyObject *current;
+ PyObject *arr;
+
+ int i, ntot, err=0;
+
+ ntot = n + nadd;
+ if (ntot < 2 || ntot > NPY_MAXARGS) {
+ PyErr_Format(PyExc_ValueError,
+ "Need between 2 and (%d) " \
+ "array objects (inclusive).", NPY_MAXARGS);
+ return NULL;
+ }
+
+ multi = _pya_malloc(sizeof(PyArrayMultiIterObject));
+ if (multi == NULL) return PyErr_NoMemory();
+ PyObject_Init((PyObject *)multi, &PyArrayMultiIter_Type);
+
+ for(i=0; i<ntot; i++) multi->iters[i] = NULL;
+ multi->numiter = ntot;
+ multi->index = 0;
+
+ va_start(va, nadd);
+ for(i=0; i<ntot; i++) {
+ if (i < n) {
+ current = mps[i];
+ }
+ else {
+ current = va_arg(va, PyObject *);
+ }
+ arr = PyArray_FROM_O(current);
+ if (arr==NULL) {
+ err=1; break;
+ }
+ else {
+ multi->iters[i] = (PyArrayIterObject *)PyArray_IterNew(arr);
+ Py_DECREF(arr);
+ }
+ }
+
+ va_end(va);
+
+ if (!err && PyArray_Broadcast(multi) < 0) err=1;
+
+ if (err) {
+ Py_DECREF(multi);
+ return NULL;
+ }
+
+ PyArray_MultiIter_RESET(multi);
+
+ return (PyObject *)multi;
+}
+
/*NUMPY_API
Get MultiIterator,
*/