Fix ticket #925

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,
 */