1
#define PY_SSIZE_T_CLEAN
2
#include <Python.h>
3
#include "structmember.h"
4

5
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
6
#define _MULTIARRAYMODULE
7
#include "numpy/arrayobject.h"
8
#include "numpy/arrayscalars.h"
9

10
#include "npy_config.h"
11

12
#include "npy_pycompat.h"
13

14
#include "arrayobject.h"
15
#include "iterators.h"
16
#include "ctors.h"
17
#include "common.h"
18
#include "array_coercion.h"
19

20
#define NEWAXIS_INDEX -1
21
#define ELLIPSIS_INDEX -2
22
#define SINGLE_INDEX -3
23

24
/*
25
 * Tries to convert 'o' into an npy_intp interpreted as an
26
 * index. Returns 1 if it was successful, 0 otherwise. Does
27
 * not set an exception.
28
 */
29
static int
30 1
coerce_index(PyObject *o, npy_intp *v)
31
{
32 1
    *v = PyArray_PyIntAsIntp(o);
33

34 1
    if ((*v) == -1 && PyErr_Occurred()) {
35 0
        PyErr_Clear();
36 0
        return 0;
37
    }
38
    return 1;
39
}
40

41
/*
42
 * This function converts one element of the indexing tuple
43
 * into a step size and a number of steps, returning the
44
 * starting index. Non-slices are signalled in 'n_steps',
45
 * as NEWAXIS_INDEX, ELLIPSIS_INDEX, or SINGLE_INDEX.
46
 */
47
NPY_NO_EXPORT npy_intp
48 1
parse_index_entry(PyObject *op, npy_intp *step_size,
49
                  npy_intp *n_steps, npy_intp max,
50
                  int axis, int check_index)
51
{
52
    npy_intp i;
53

54 1
    if (op == Py_None) {
55 0
        *n_steps = NEWAXIS_INDEX;
56 0
        i = 0;
57
    }
58 1
    else if (op == Py_Ellipsis) {
59 0
        *n_steps = ELLIPSIS_INDEX;
60 0
        i = 0;
61
    }
62 1
    else if (PySlice_Check(op)) {
63
        npy_intp stop;
64 1
        if (PySlice_GetIndicesEx(op, max, &i, &stop, step_size, n_steps) < 0) {
65
            goto fail;
66
        }
67 1
        if (*n_steps <= 0) {
68 1
            *n_steps = 0;
69 1
            *step_size = 1;
70 1
            i = 0;
71
        }
72
    }
73 1
    else if (coerce_index(op, &i)) {
74 1
        *n_steps = SINGLE_INDEX;
75 1
        *step_size = 0;
76 1
        if (check_index) {
77 1
            if (check_and_adjust_index(&i, max, axis, NULL) < 0) {
78
                goto fail;
79
            }
80
        }
81
    }
82
    else {
83 0
        PyErr_SetString(PyExc_IndexError,
84
                        "each index entry must be either a "
85
                        "slice, an integer, Ellipsis, or "
86
                        "newaxis");
87 0
        goto fail;
88
    }
89 1
    return i;
90

91 1
 fail:
92
    return -1;
93
}
94

95

96
/*********************** Element-wise Array Iterator ***********************/
97
/*  Aided by Peter J. Verveer's  nd_image package and numpy's arraymap  ****/
98
/*         and Python's array iterator                                   ***/
99

100
/* get the dataptr from its current coordinates for simple iterator */
101
static char*
102 1
get_ptr_simple(PyArrayIterObject* iter, const npy_intp *coordinates)
103
{
104
    npy_intp i;
105
    char *ret;
106

107 1
    ret = PyArray_DATA(iter->ao);
108

109 1
    for(i = 0; i < PyArray_NDIM(iter->ao); ++i) {
110 1
            ret += coordinates[i] * iter->strides[i];
111
    }
112

113 1
    return ret;
114
}
115

116
/*
117
 * This is common initialization code between PyArrayIterObject and
118
 * PyArrayNeighborhoodIterObject
119
 *
120
 * Steals a reference to the array object which gets removed at deallocation,
121
 * if the iterator is allocated statically and its dealloc not called, it
122
 * can be thought of as borrowing the reference.
123
 */
124
NPY_NO_EXPORT void
125 1
PyArray_RawIterBaseInit(PyArrayIterObject *it, PyArrayObject *ao)
126
{
127
    int nd, i;
128

129 1
    nd = PyArray_NDIM(ao);
130 1
    PyArray_UpdateFlags(ao, NPY_ARRAY_C_CONTIGUOUS);
131 1
    if (PyArray_ISCONTIGUOUS(ao)) {
132 1
        it->contiguous = 1;
133
    }
134
    else {
135 1
        it->contiguous = 0;
136
    }
137 1
    it->ao = ao;
138 1
    it->size = PyArray_SIZE(ao);
139 1
    it->nd_m1 = nd - 1;
140 1
    if (nd != 0) {
141 1
        it->factors[nd-1] = 1;
142
    }
143 1
    for (i = 0; i < nd; i++) {
144 1
        it->dims_m1[i] = PyArray_DIMS(ao)[i] - 1;
145 1
        it->strides[i] = PyArray_STRIDES(ao)[i];
146 1
        it->backstrides[i] = it->strides[i] * it->dims_m1[i];
147 1
        if (i > 0) {
148 1
            it->factors[nd-i-1] = it->factors[nd-i] * PyArray_DIMS(ao)[nd-i];
149
        }
150 1
        it->bounds[i][0] = 0;
151 1
        it->bounds[i][1] = PyArray_DIMS(ao)[i] - 1;
152 1
        it->limits[i][0] = 0;
153 1
        it->limits[i][1] = PyArray_DIMS(ao)[i] - 1;
154 1
        it->limits_sizes[i] = it->limits[i][1] - it->limits[i][0] + 1;
155
    }
156

157 1
    it->translate = &get_ptr_simple;
158 1
    PyArray_ITER_RESET(it);
159

160 1
    return;
161
}
162

163
static void
164
array_iter_base_dealloc(PyArrayIterObject *it)
165
{
166 1
    Py_XDECREF(it->ao);
167
}
168

169
/*NUMPY_API
170
 * Get Iterator.
171
 */
172
NPY_NO_EXPORT PyObject *
173 1
PyArray_IterNew(PyObject *obj)
174
{
175
    /*
176
     * Note that internally PyArray_RawIterBaseInit may be called directly on a
177
     * statically allocated PyArrayIterObject.
178
     */
179
    PyArrayIterObject *it;
180
    PyArrayObject *ao;
181

182 1
    if (!PyArray_Check(obj)) {
183 0
        PyErr_BadInternalCall();
184 0
        return NULL;
185
    }
186 1
    ao = (PyArrayObject *)obj;
187

188 1
    it = (PyArrayIterObject *)PyArray_malloc(sizeof(PyArrayIterObject));
189 1
    PyObject_Init((PyObject *)it, &PyArrayIter_Type);
190
    /* it = PyObject_New(PyArrayIterObject, &PyArrayIter_Type);*/
191 1
    if (it == NULL) {
192
        return NULL;
193
    }
194

195 1
    Py_INCREF(ao);  /* PyArray_RawIterBaseInit steals a reference */
196 1
    PyArray_RawIterBaseInit(it, ao);
197 1
    return (PyObject *)it;
198
}
199

200
/*NUMPY_API
201
 * Get Iterator broadcast to a particular shape
202
 */
203
NPY_NO_EXPORT PyObject *
204 1
PyArray_BroadcastToShape(PyObject *obj, npy_intp *dims, int nd)
205
{
206
    PyArrayIterObject *it;
207
    int i, diff, j, compat, k;
208 1
    PyArrayObject *ao = (PyArrayObject *)obj;
209

210 1
    if (PyArray_NDIM(ao) > nd) {
211
        goto err;
212
    }
213 1
    compat = 1;
214 1
    diff = j = nd - PyArray_NDIM(ao);
215 1
    for (i = 0; i < PyArray_NDIM(ao); i++, j++) {
216 1
        if (PyArray_DIMS(ao)[i] == 1) {
217 1
            continue;
218
        }
219 1
        if (PyArray_DIMS(ao)[i] != dims[j]) {
220
            compat = 0;
221
            break;
222
        }
223
    }
224 1
    if (!compat) {
225
        goto err;
226
    }
227 1
    it = (PyArrayIterObject *)PyArray_malloc(sizeof(PyArrayIterObject));
228 1
    if (it == NULL) {
229
        return NULL;
230
    }
231 1
    PyObject_Init((PyObject *)it, &PyArrayIter_Type);
232

233 1
    PyArray_UpdateFlags(ao, NPY_ARRAY_C_CONTIGUOUS);
234 1
    if (PyArray_ISCONTIGUOUS(ao)) {
235 1
        it->contiguous = 1;
236
    }
237
    else {
238 1
        it->contiguous = 0;
239
    }
240 1
    Py_INCREF(ao);
241 1
    it->ao = ao;
242 1
    it->size = PyArray_MultiplyList(dims, nd);
243 1
    it->nd_m1 = nd - 1;
244 1
    if (nd != 0) {
245 1
        it->factors[nd-1] = 1;
246
    }
247 1
    for (i = 0; i < nd; i++) {
248 1
        it->dims_m1[i] = dims[i] - 1;
249 1
        k = i - diff;
250 1
        if ((k < 0) || PyArray_DIMS(ao)[k] != dims[i]) {
251 1
            it->contiguous = 0;
252 1
            it->strides[i] = 0;
253
        }
254
        else {
255 1
            it->strides[i] = PyArray_STRIDES(ao)[k];
256
        }
257 1
        it->backstrides[i] = it->strides[i] * it->dims_m1[i];
258 1
        if (i > 0) {
259 1
            it->factors[nd-i-1] = it->factors[nd-i] * dims[nd-i];
260
        }
261
    }
262 1
    PyArray_ITER_RESET(it);
263 1
    return (PyObject *)it;
264

265 0
 err:
266 0
    PyErr_SetString(PyExc_ValueError, "array is not broadcastable to "\
267
                    "correct shape");
268 0
    return NULL;
269
}
270

271

272

273

274

275
/*NUMPY_API
276
 * Get Iterator that iterates over all but one axis (don't use this with
277
 * PyArray_ITER_GOTO1D).  The axis will be over-written if negative
278
 * with the axis having the smallest stride.
279
 */
280
NPY_NO_EXPORT PyObject *
281 1
PyArray_IterAllButAxis(PyObject *obj, int *inaxis)
282
{
283
    PyArrayObject *arr;
284
    PyArrayIterObject *it;
285
    int axis;
286

287 1
    if (!PyArray_Check(obj)) {
288 0
        PyErr_SetString(PyExc_ValueError,
289
                "Numpy IterAllButAxis requires an ndarray");
290 0
        return NULL;
291
    }
292 1
    arr = (PyArrayObject *)obj;
293

294 1
    it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)arr);
295 1
    if (it == NULL) {
296
        return NULL;
297
    }
298 1
    if (PyArray_NDIM(arr)==0) {
299
        return (PyObject *)it;
300
    }
301 1
    if (*inaxis < 0) {
302
        int i, minaxis = 0;
303
        npy_intp minstride = 0;
304
        i = 0;
305 0
        while (minstride == 0 && i < PyArray_NDIM(arr)) {
306 0
            minstride = PyArray_STRIDE(arr,i);
307 0
            i++;
308
        }
309 0
        for (i = 1; i < PyArray_NDIM(arr); i++) {
310 0
            if (PyArray_STRIDE(arr,i) > 0 &&
311 0
                PyArray_STRIDE(arr, i) < minstride) {
312 0
                minaxis = i;
313 0
                minstride = PyArray_STRIDE(arr,i);
314
            }
315
        }
316 0
        *inaxis = minaxis;
317
    }
318 1
    axis = *inaxis;
319
    /* adjust so that will not iterate over axis */
320 1
    it->contiguous = 0;
321 1
    if (it->size != 0) {
322 1
        it->size /= PyArray_DIM(arr,axis);
323
    }
324 1
    it->dims_m1[axis] = 0;
325 1
    it->backstrides[axis] = 0;
326

327
    /*
328
     * (won't fix factors so don't use
329
     * PyArray_ITER_GOTO1D with this iterator)
330
     */
331 1
    return (PyObject *)it;
332
}
333

334
/*NUMPY_API
335
 * Adjusts previously broadcasted iterators so that the axis with
336
 * the smallest sum of iterator strides is not iterated over.
337
 * Returns dimension which is smallest in the range [0,multi->nd).
338
 * A -1 is returned if multi->nd == 0.
339
 *
340
 * don't use with PyArray_ITER_GOTO1D because factors are not adjusted
341
 */
342
NPY_NO_EXPORT int
343 0
PyArray_RemoveSmallest(PyArrayMultiIterObject *multi)
344
{
345
    PyArrayIterObject *it;
346
    int i, j;
347
    int axis;
348
    npy_intp smallest;
349
    npy_intp sumstrides[NPY_MAXDIMS];
350

351 0
    if (multi->nd == 0) {
352
        return -1;
353
    }
354 0
    for (i = 0; i < multi->nd; i++) {
355 0
        sumstrides[i] = 0;
356 0
        for (j = 0; j < multi->numiter; j++) {
357 0
            sumstrides[i] += multi->iters[j]->strides[i];
358
        }
359
    }
360 0
    axis = 0;
361 0
    smallest = sumstrides[0];
362
    /* Find longest dimension */
363 0
    for (i = 1; i < multi->nd; i++) {
364 0
        if (sumstrides[i] < smallest) {
365 0
            axis = i;
366 0
            smallest = sumstrides[i];
367
        }
368
    }
369 0
    for(i = 0; i < multi->numiter; i++) {
370 0
        it = multi->iters[i];
371 0
        it->contiguous = 0;
372 0
        if (it->size != 0) {
373 0
            it->size /= (it->dims_m1[axis]+1);
374
        }
375 0
        it->dims_m1[axis] = 0;
376 0
        it->backstrides[axis] = 0;
377
    }
378 0
    multi->size = multi->iters[0]->size;
379 0
    return axis;
380
}
381

382
/* Returns an array scalar holding the element desired */
383

384
static PyObject *
385 1
arrayiter_next(PyArrayIterObject *it)
386
{
387
    PyObject *ret;
388

389 1
    if (it->index < it->size) {
390 1
        ret = PyArray_ToScalar(it->dataptr, it->ao);
391 1
        PyArray_ITER_NEXT(it);
392
        return ret;
393
    }
394
    return NULL;
395
}
396

397
static void
398 1
arrayiter_dealloc(PyArrayIterObject *it)
399
{
400
    /*
401
     * Note that it is possible to statically allocate a PyArrayIterObject,
402
     * which does not call this function.
403
     */
404 1
    array_iter_base_dealloc(it);
405 1
    PyArray_free(it);
406
}
407

408
static Py_ssize_t
409 1
iter_length(PyArrayIterObject *self)
410
{
411 1
    return self->size;
412
}
413

414

415
static PyArrayObject *
416 1
iter_subscript_Bool(PyArrayIterObject *self, PyArrayObject *ind)
417
{
418
    npy_intp counter, strides;
419
    int itemsize;
420 1
    npy_intp count = 0;
421
    char *dptr, *optr;
422
    PyArrayObject *ret;
423
    int swap;
424
    PyArray_CopySwapFunc *copyswap;
425

426

427 1
    if (PyArray_NDIM(ind) != 1) {
428 0
        PyErr_SetString(PyExc_ValueError,
429
                        "boolean index array should have 1 dimension");
430 0
        return NULL;
431
    }
432 1
    counter = PyArray_DIMS(ind)[0];
433 1
    if (counter > self->size) {
434 0
        PyErr_SetString(PyExc_ValueError,
435
                        "too many boolean indices");
436 0
        return NULL;
437
    }
438

439 1
    strides = PyArray_STRIDES(ind)[0];
440 1
    dptr = PyArray_DATA(ind);
441
    /* Get size of return array */
442 1
    while (counter--) {
443 1
        if (*((npy_bool *)dptr) != 0) {
444 1
            count++;
445
        }
446 1
        dptr += strides;
447
    }
448 1
    itemsize = PyArray_DESCR(self->ao)->elsize;
449 1
    Py_INCREF(PyArray_DESCR(self->ao));
450 1
    ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self->ao),
451
                             PyArray_DESCR(self->ao), 1, &count,
452
                             NULL, NULL,
453
                             0, (PyObject *)self->ao);
454 1
    if (ret == NULL) {
455
        return NULL;
456
    }
457
    /* Set up loop */
458 1
    optr = PyArray_DATA(ret);
459 1
    counter = PyArray_DIMS(ind)[0];
460 1
    dptr = PyArray_DATA(ind);
461 1
    copyswap = PyArray_DESCR(self->ao)->f->copyswap;
462
    /* Loop over Boolean array */
463 1
    swap = (PyArray_ISNOTSWAPPED(self->ao) != PyArray_ISNOTSWAPPED(ret));
464 1
    while (counter--) {
465 1
        if (*((npy_bool *)dptr) != 0) {
466 1
            copyswap(optr, self->dataptr, swap, self->ao);
467 1
            optr += itemsize;
468
        }
469 1
        dptr += strides;
470 1
        PyArray_ITER_NEXT(self);
471
    }
472 1
    PyArray_ITER_RESET(self);
473 1
    return ret;
474
}
475

476
static PyObject *
477 1
iter_subscript_int(PyArrayIterObject *self, PyArrayObject *ind)
478
{
479
    npy_intp num;
480
    PyArrayObject *ret;
481
    PyArrayIterObject *ind_it;
482
    int itemsize;
483
    int swap;
484
    char *optr;
485
    npy_intp counter;
486
    PyArray_CopySwapFunc *copyswap;
487

488 1
    itemsize = PyArray_DESCR(self->ao)->elsize;
489 1
    if (PyArray_NDIM(ind) == 0) {
490 1
        num = *((npy_intp *)PyArray_DATA(ind));
491 1
        if (check_and_adjust_index(&num, self->size, -1, NULL) < 0) {
492 1
            PyArray_ITER_RESET(self);
493 1
            return NULL;
494
        }
495
        else {
496
            PyObject *tmp;
497 1
            PyArray_ITER_GOTO1D(self, num);
498 1
            tmp = PyArray_ToScalar(self->dataptr, self->ao);
499 1
            PyArray_ITER_RESET(self);
500 1
            return tmp;
501
        }
502
    }
503

504 1
    Py_INCREF(PyArray_DESCR(self->ao));
505 1
    ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self->ao),
506
                             PyArray_DESCR(self->ao),
507
                             PyArray_NDIM(ind),
508 1
                             PyArray_DIMS(ind),
509
                             NULL, NULL,
510 1
                             0, (PyObject *)self->ao);
511 1
    if (ret == NULL) {
512
        return NULL;
513
    }
514 1
    optr = PyArray_DATA(ret);
515 1
    ind_it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)ind);
516 1
    if (ind_it == NULL) {
517 0
        Py_DECREF(ret);
518
        return NULL;
519
    }
520 1
    counter = ind_it->size;
521 1
    copyswap = PyArray_DESCR(ret)->f->copyswap;
522 1
    swap = (PyArray_ISNOTSWAPPED(ret) != PyArray_ISNOTSWAPPED(self->ao));
523 1
    while (counter--) {
524 1
        num = *((npy_intp *)(ind_it->dataptr));
525 1
        if (check_and_adjust_index(&num, self->size, -1, NULL) < 0) {
526 1
            Py_DECREF(ind_it);
527 1
            Py_DECREF(ret);
528 1
            PyArray_ITER_RESET(self);
529 1
            return NULL;
530
        }
531 1
        PyArray_ITER_GOTO1D(self, num);
532 1
        copyswap(optr, self->dataptr, swap, ret);
533 1
        optr += itemsize;
534 1
        PyArray_ITER_NEXT(ind_it);
535
    }
536 1
    Py_DECREF(ind_it);
537 1
    PyArray_ITER_RESET(self);
538 1
    return (PyObject *)ret;
539
}
540

541
/* Always returns arrays */
542
NPY_NO_EXPORT PyObject *
543 1
iter_subscript(PyArrayIterObject *self, PyObject *ind)
544
{
545 1
    PyArray_Descr *indtype = NULL;
546
    PyArray_Descr *dtype;
547
    npy_intp start, step_size;
548
    npy_intp n_steps;
549
    PyArrayObject *ret;
550
    char *dptr;
551
    int size;
552 1
    PyObject *obj = NULL;
553
    PyObject *new;
554
    PyArray_CopySwapFunc *copyswap;
555

556 1
    if (ind == Py_Ellipsis) {
557 1
        ind = PySlice_New(NULL, NULL, NULL);
558 1
        obj = iter_subscript(self, ind);
559 1
        Py_DECREF(ind);
560
        return obj;
561
    }
562 1
    if (PyTuple_Check(ind)) {
563
        int len;
564 1
        len = PyTuple_GET_SIZE(ind);
565 1
        if (len > 1) {
566
            goto fail;
567
        }
568 1
        if (len == 0) {
569 0
            Py_INCREF(self->ao);
570 0
            return (PyObject *)self->ao;
571
        }
572 1
        ind = PyTuple_GET_ITEM(ind, 0);
573
    }
574

575
    /*
576
     * Tuples >1d not accepted --- i.e. no newaxis
577
     * Could implement this with adjusted strides and dimensions in iterator
578
     * Check for Boolean -- this is first because Bool is a subclass of Int
579
     */
580 1
    PyArray_ITER_RESET(self);
581

582 1
    if (PyBool_Check(ind)) {
583 0
        if (PyObject_IsTrue(ind)) {
584 0
            return PyArray_ToScalar(self->dataptr, self->ao);
585
        }
586
        else { /* empty array */
587 0
            npy_intp ii = 0;
588 0
            dtype = PyArray_DESCR(self->ao);
589 0
            Py_INCREF(dtype);
590 0
            ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self->ao),
591
                                     dtype,
592
                                     1, &ii,
593
                                     NULL, NULL, 0,
594
                                     (PyObject *)self->ao);
595
            return (PyObject *)ret;
596
        }
597
    }
598

599
    /* Check for Integer or Slice */
600 1
    if (PyLong_Check(ind) || PyInt_Check(ind) || PySlice_Check(ind)) {
601 1
        start = parse_index_entry(ind, &step_size, &n_steps,
602
                                  self->size, 0, 1);
603 1
        if (start == -1) {
604
            goto fail;
605
        }
606 1
        if (n_steps == ELLIPSIS_INDEX || n_steps == NEWAXIS_INDEX) {
607 0
            PyErr_SetString(PyExc_IndexError,
608
                            "cannot use Ellipsis or newaxes here");
609 0
            goto fail;
610
        }
611 1
        PyArray_ITER_GOTO1D(self, start);
612 1
        if (n_steps == SINGLE_INDEX) { /* Integer */
613
            PyObject *tmp;
614 1
            tmp = PyArray_ToScalar(self->dataptr, self->ao);
615 1
            PyArray_ITER_RESET(self);
616 1
            return tmp;
617
        }
618 1
        size = PyArray_DESCR(self->ao)->elsize;
619 1
        dtype = PyArray_DESCR(self->ao);
620 1
        Py_INCREF(dtype);
621 1
        ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self->ao),
622
                                 dtype,
623
                                 1, &n_steps,
624
                                 NULL, NULL,
625
                                 0, (PyObject *)self->ao);
626 1
        if (ret == NULL) {
627
            goto fail;
628
        }
629 1
        dptr = PyArray_DATA(ret);
630 1
        copyswap = PyArray_DESCR(ret)->f->copyswap;
631 1
        while (n_steps--) {
632 1
            copyswap(dptr, self->dataptr, 0, ret);
633 1
            start += step_size;
634 1
            PyArray_ITER_GOTO1D(self, start);
635 1
            dptr += size;
636
        }
637 1
        PyArray_ITER_RESET(self);
638 1
        return (PyObject *)ret;
639
    }
640

641
    /* convert to INTP array if Integer array scalar or List */
642 1
    indtype = PyArray_DescrFromType(NPY_INTP);
643 1
    if (PyArray_IsScalar(ind, Integer) || PyList_Check(ind)) {
644 1
        Py_INCREF(indtype);
645 1
        obj = PyArray_FromAny(ind, indtype, 0, 0, NPY_ARRAY_FORCECAST, NULL);
646 1
        if (obj == NULL) {
647
            goto fail;
648
        }
649
    }
650
    else {
651 1
        Py_INCREF(ind);
652 1
        obj = ind;
653
    }
654

655
    /* Any remaining valid input is an array or has been turned into one */
656 1
    if (!PyArray_Check(obj)) {
657
        goto fail;
658
    }
659

660
    /* Check for Boolean array */
661 1
    if (PyArray_TYPE((PyArrayObject *)obj) == NPY_BOOL) {
662 1
        ret = iter_subscript_Bool(self, (PyArrayObject *)obj);
663 1
        Py_DECREF(indtype);
664 1
        Py_DECREF(obj);
665
        return (PyObject *)ret;
666
    }
667

668
    /* Only integer arrays left */
669 1
    if (!PyArray_ISINTEGER((PyArrayObject *)obj)) {
670
        goto fail;
671
    }
672

673 1
    Py_INCREF(indtype);
674 1
    new = PyArray_FromAny(obj, indtype, 0, 0,
675
                      NPY_ARRAY_FORCECAST | NPY_ARRAY_ALIGNED, NULL);
676 1
    if (new == NULL) {
677
        goto fail;
678
    }
679 1
    Py_DECREF(indtype);
680 1
    Py_DECREF(obj);
681 1
    ret = (PyArrayObject *)iter_subscript_int(self, (PyArrayObject *)new);
682 1
    Py_DECREF(new);
683
    return (PyObject *)ret;
684

685

686 1
 fail:
687 1
    if (!PyErr_Occurred()) {
688 1
        PyErr_SetString(PyExc_IndexError, "unsupported iterator index");
689
    }
690 1
    Py_XDECREF(indtype);
691 1
    Py_XDECREF(obj);
692
    return NULL;
693

694
}
695

696

697
static int
698 1
iter_ass_sub_Bool(PyArrayIterObject *self, PyArrayObject *ind,
699
                  PyArrayIterObject *val, int swap)
700
{
701
    npy_intp counter, strides;
702
    char *dptr;
703
    PyArray_CopySwapFunc *copyswap;
704

705 1
    if (PyArray_NDIM(ind) != 1) {
706 0
        PyErr_SetString(PyExc_ValueError,
707
                        "boolean index array should have 1 dimension");
708 0
        return -1;
709
    }
710

711 1
    counter = PyArray_DIMS(ind)[0];
712 1
    if (counter > self->size) {
713 1
        PyErr_SetString(PyExc_ValueError,
714
                        "boolean index array has too many values");
715 1
        return -1;
716
    }
717

718 0
    strides = PyArray_STRIDES(ind)[0];
719 0
    dptr = PyArray_DATA(ind);
720 0
    PyArray_ITER_RESET(self);
721
    /* Loop over Boolean array */
722 0
    copyswap = PyArray_DESCR(self->ao)->f->copyswap;
723 0
    while (counter--) {
724 0
        if (*((npy_bool *)dptr) != 0) {
725 0
            copyswap(self->dataptr, val->dataptr, swap, self->ao);
726 0
            PyArray_ITER_NEXT(val);
727 0
            if (val->index == val->size) {
728 0
                PyArray_ITER_RESET(val);
729
            }
730
        }
731 0
        dptr += strides;
732 0
        PyArray_ITER_NEXT(self);
733
    }
734 0
    PyArray_ITER_RESET(self);
735 0
    return 0;
736
}
737

738
static int
739 1
iter_ass_sub_int(PyArrayIterObject *self, PyArrayObject *ind,
740
                 PyArrayIterObject *val, int swap)
741
{
742
    npy_intp num;
743
    PyArrayIterObject *ind_it;
744
    npy_intp counter;
745
    PyArray_CopySwapFunc *copyswap;
746

747 1
    copyswap = PyArray_DESCR(self->ao)->f->copyswap;
748 1
    if (PyArray_NDIM(ind) == 0) {
749 1
        num = *((npy_intp *)PyArray_DATA(ind));
750 1
        if (check_and_adjust_index(&num, self->size, -1, NULL) < 0) {
751
            return -1;
752
        }
753 0
        PyArray_ITER_GOTO1D(self, num);
754 0
        copyswap(self->dataptr, val->dataptr, swap, self->ao);
755 0
        return 0;
756
    }
757 1
    ind_it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)ind);
758 1
    if (ind_it == NULL) {
759
        return -1;
760
    }
761 1
    counter = ind_it->size;
762 1
    while (counter--) {
763 1
        num = *((npy_intp *)(ind_it->dataptr));
764 1
        if (check_and_adjust_index(&num, self->size, -1, NULL) < 0) {
765 1
            Py_DECREF(ind_it);
766
            return -1;
767
        }
768 1
        PyArray_ITER_GOTO1D(self, num);
769 1
        copyswap(self->dataptr, val->dataptr, swap, self->ao);
770 1
        PyArray_ITER_NEXT(ind_it);
771 1
        PyArray_ITER_NEXT(val);
772 1
        if (val->index == val->size) {
773 1
            PyArray_ITER_RESET(val);
774
        }
775
    }
776 1
    Py_DECREF(ind_it);
777
    return 0;
778
}
779

780
NPY_NO_EXPORT int
781 1
iter_ass_subscript(PyArrayIterObject *self, PyObject *ind, PyObject *val)
782
{
783 1
    PyArrayObject *arrval = NULL;
784 1
    PyArrayIterObject *val_it = NULL;
785
    PyArray_Descr *type;
786 1
    PyArray_Descr *indtype = NULL;
787 1
    int swap, retval = -1;
788
    npy_intp start, step_size;
789
    npy_intp n_steps;
790 1
    PyObject *obj = NULL;
791
    PyArray_CopySwapFunc *copyswap;
792

793

794 1
    if (val == NULL) {
795 1
        PyErr_SetString(PyExc_TypeError,
796
                "Cannot delete iterator elements");
797 1
        return -1;
798
    }
799

800 1
    if (PyArray_FailUnlessWriteable(self->ao, "underlying array") < 0)
801
        return -1;
802

803 1
    if (ind == Py_Ellipsis) {
804 0
        ind = PySlice_New(NULL, NULL, NULL);
805 0
        retval = iter_ass_subscript(self, ind, val);
806 0
        Py_DECREF(ind);
807
        return retval;
808
    }
809

810 1
    if (PyTuple_Check(ind)) {
811
        int len;
812 1
        len = PyTuple_GET_SIZE(ind);
813 1
        if (len > 1) {
814
            goto finish;
815
        }
816 1
        ind = PyTuple_GET_ITEM(ind, 0);
817
    }
818

819 1
    type = PyArray_DESCR(self->ao);
820

821
    /*
822
     * Check for Boolean -- this is first because
823
     * Bool is a subclass of Int
824
     */
825 1
    if (PyBool_Check(ind)) {
826 0
        retval = 0;
827 0
        if (PyObject_IsTrue(ind)) {
828 0
            retval = PyArray_Pack(PyArray_DESCR(self->ao), self->dataptr, val);
829
        }
830
        goto finish;
831
    }
832

833 1
    if (PySequence_Check(ind) || PySlice_Check(ind)) {
834
        goto skip;
835
    }
836 1
    start = PyArray_PyIntAsIntp(ind);
837 1
    if (error_converting(start)) {
838 0
        PyErr_Clear();
839
    }
840
    else {
841 1
        if (check_and_adjust_index(&start, self->size, -1, NULL) < 0) {
842
            goto finish;
843
        }
844 1
        PyArray_ITER_GOTO1D(self, start);
845 1
        retval = PyArray_Pack(PyArray_DESCR(self->ao), self->dataptr, val);
846 1
        PyArray_ITER_RESET(self);
847 1
        if (retval < 0) {
848 0
            PyErr_SetString(PyExc_ValueError,
849
                            "Error setting single item of array.");
850
        }
851
        goto finish;
852
    }
853

854 1
 skip:
855 1
    Py_INCREF(type);
856 1
    arrval = (PyArrayObject *)PyArray_FromAny(val, type, 0, 0,
857
                                              NPY_ARRAY_FORCECAST, NULL);
858 1
    if (arrval == NULL) {
859
        return -1;
860
    }
861 1
    val_it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)arrval);
862 1
    if (val_it == NULL) {
863
        goto finish;
864
    }
865 1
    if (val_it->size == 0) {
866
        retval = 0;
867
        goto finish;
868
    }
869

870 1
    copyswap = PyArray_DESCR(arrval)->f->copyswap;
871 1
    swap = (PyArray_ISNOTSWAPPED(self->ao)!=PyArray_ISNOTSWAPPED(arrval));
872

873
    /* Check Slice */
874 1
    if (PySlice_Check(ind)) {
875 1
        start = parse_index_entry(ind, &step_size, &n_steps, self->size, 0, 0);
876 1
        if (start == -1) {
877
            goto finish;
878
        }
879 1
        if (n_steps == ELLIPSIS_INDEX || n_steps == NEWAXIS_INDEX) {
880 0
            PyErr_SetString(PyExc_IndexError,
881
                            "cannot use Ellipsis or newaxes here");
882 0
            goto finish;
883
        }
884 1
        PyArray_ITER_GOTO1D(self, start);
885 1
        if (n_steps == SINGLE_INDEX) {
886
            /* Integer */
887 0
            copyswap(self->dataptr, PyArray_DATA(arrval), swap, arrval);
888 0
            PyArray_ITER_RESET(self);
889 0
            retval = 0;
890 0
            goto finish;
891
        }
892 1
        while (n_steps--) {
893 1
            copyswap(self->dataptr, val_it->dataptr, swap, arrval);
894 1
            start += step_size;
895 1
            PyArray_ITER_GOTO1D(self, start);
896 1
            PyArray_ITER_NEXT(val_it);
897 1
            if (val_it->index == val_it->size) {
898 1
                PyArray_ITER_RESET(val_it);
899
            }
900
        }
901 1
        PyArray_ITER_RESET(self);
902 1
        retval = 0;
903 1
        goto finish;
904
    }
905

906
    /* convert to INTP array if Integer array scalar or List */
907 1
    indtype = PyArray_DescrFromType(NPY_INTP);
908 1
    if (PyList_Check(ind)) {
909 0
        Py_INCREF(indtype);
910 0
        obj = PyArray_FromAny(ind, indtype, 0, 0, NPY_ARRAY_FORCECAST, NULL);
911
    }
912
    else {
913 1
        Py_INCREF(ind);
914 1
        obj = ind;
915
    }
916

917 1
    if (obj != NULL && PyArray_Check(obj)) {
918
        /* Check for Boolean object */
919 1
        if (PyArray_TYPE((PyArrayObject *)obj)==NPY_BOOL) {
920 1
            if (iter_ass_sub_Bool(self, (PyArrayObject *)obj,
921
                                  val_it, swap) < 0) {
922
                goto finish;
923
            }
924 0
            retval=0;
925
        }
926
        /* Check for integer array */
927 1
        else if (PyArray_ISINTEGER((PyArrayObject *)obj)) {
928
            PyObject *new;
929 1
            Py_INCREF(indtype);
930 1
            new = PyArray_CheckFromAny(obj, indtype, 0, 0,
931
                           NPY_ARRAY_FORCECAST | NPY_ARRAY_BEHAVED_NS, NULL);
932 1
            Py_DECREF(obj);
933 1
            obj = new;
934 1
            if (new == NULL) {
935
                goto finish;
936
            }
937 1
            if (iter_ass_sub_int(self, (PyArrayObject *)obj,
938
                                 val_it, swap) < 0) {
939
                goto finish;
940
            }
941 1
            retval = 0;
942
        }
943
    }
944

945 1
 finish:
946 1
    if (!PyErr_Occurred() && retval < 0) {
947 0
        PyErr_SetString(PyExc_IndexError, "unsupported iterator index");
948
    }
949 1
    Py_XDECREF(indtype);
950 1
    Py_XDECREF(obj);
951 1
    Py_XDECREF(val_it);
952 1
    Py_XDECREF(arrval);
953
    return retval;
954

955
}
956

957

958
static PyMappingMethods iter_as_mapping = {
959
    (lenfunc)iter_length,                   /*mp_length*/
960
    (binaryfunc)iter_subscript,             /*mp_subscript*/
961
    (objobjargproc)iter_ass_subscript,      /*mp_ass_subscript*/
962
};
963

964

965
/* Two options:
966
 *  1) underlying array is contiguous
967
 *     -- return 1-d wrapper around it
968
 *  2) underlying array is not contiguous
969
 *     -- make new 1-d contiguous array with updateifcopy flag set
970
 *        to copy back to the old array
971
 *
972
 *  If underlying array is readonly, then we make the output array readonly
973
 *     and updateifcopy does not apply.
974
 *
975
 *  Changed 2017-07-21, 1.14.0.
976
 *
977
 *  In order to start the process of removing UPDATEIFCOPY, see gh-7054, the
978
 *  behavior is changed to always return an non-writeable copy when the base
979
 *  array is non-contiguous. Doing that will hopefully smoke out those few
980
 *  folks who assign to the result with the expectation that the base array
981
 *  will be changed. At a later date non-contiguous arrays will always return
982
 *  writeable copies.
983
 *
984
 *  Note that the type and argument expected for the __array__ method is
985
 *  ignored.
986
 */
987
static PyArrayObject *
988 1
iter_array(PyArrayIterObject *it, PyObject *NPY_UNUSED(op))
989
{
990

991
    PyArrayObject *ret;
992
    npy_intp size;
993

994 1
    size = PyArray_SIZE(it->ao);
995 1
    Py_INCREF(PyArray_DESCR(it->ao));
996

997 1
    if (PyArray_ISCONTIGUOUS(it->ao)) {
998 1
        ret = (PyArrayObject *)PyArray_NewFromDescrAndBase(
999
                &PyArray_Type, PyArray_DESCR(it->ao),
1000
                1, &size, NULL, PyArray_DATA(it->ao),
1001 1
                PyArray_FLAGS(it->ao), (PyObject *)it->ao, (PyObject *)it->ao);
1002 1
        if (ret == NULL) {
1003
            return NULL;
1004
        }
1005
    }
1006
    else {
1007 1
        ret = (PyArrayObject *)PyArray_NewFromDescr(
1008
                &PyArray_Type, PyArray_DESCR(it->ao), 1, &size,
1009
                NULL, NULL, 0,
1010
                (PyObject *)it->ao);
1011 1
        if (ret == NULL) {
1012
            return NULL;
1013
        }
1014 1
        if (PyArray_CopyAnyInto(ret, it->ao) < 0) {
1015 0
            Py_DECREF(ret);
1016
            return NULL;
1017
        }
1018
        PyArray_CLEARFLAGS(ret, NPY_ARRAY_WRITEABLE);
1019
    }
1020
    return ret;
1021

1022
}
1023

1024
static PyObject *
1025 1
iter_copy(PyArrayIterObject *it, PyObject *args)
1026
{
1027 1
    if (!PyArg_ParseTuple(args, "")) {
1028
        return NULL;
1029
    }
1030 1
    return PyArray_Flatten(it->ao, 0);
1031
}
1032

1033
static PyMethodDef iter_methods[] = {
1034
    /* to get array */
1035
    {"__array__",
1036
        (PyCFunction)iter_array,
1037
        METH_VARARGS, NULL},
1038
    {"copy",
1039
        (PyCFunction)iter_copy,
1040
        METH_VARARGS, NULL},
1041
    {NULL, NULL, 0, NULL}           /* sentinel */
1042
};
1043

1044
static PyObject *
1045 0
iter_richcompare(PyArrayIterObject *self, PyObject *other, int cmp_op)
1046
{
1047
    PyArrayObject *new;
1048
    PyObject *ret;
1049 0
    new = (PyArrayObject *)iter_array(self, NULL);
1050 0
    if (new == NULL) {
1051
        return NULL;
1052
    }
1053 0
    ret = array_richcompare(new, other, cmp_op);
1054 0
    PyArray_ResolveWritebackIfCopy(new);
1055 0
    Py_DECREF(new);
1056
    return ret;
1057
}
1058

1059

1060
static PyMemberDef iter_members[] = {
1061
    {"base",
1062
        T_OBJECT,
1063
        offsetof(PyArrayIterObject, ao),
1064
        READONLY, NULL},
1065
    {"index",
1066
        T_INT,
1067
        offsetof(PyArrayIterObject, index),
1068
        READONLY, NULL},
1069
    {NULL, 0, 0, 0, NULL},
1070
};
1071

1072
static PyObject *
1073 1
iter_coords_get(PyArrayIterObject *self)
1074
{
1075
    int nd;
1076 1
    nd = PyArray_NDIM(self->ao);
1077 1
    if (self->contiguous) {
1078
        /*
1079
         * coordinates not kept track of ---
1080
         * need to generate from index
1081
         */
1082
        npy_intp val;
1083
        int i;
1084 1
        val = self->index;
1085 1
        for (i = 0; i < nd; i++) {
1086 1
            if (self->factors[i] != 0) {
1087 1
                self->coordinates[i] = val / self->factors[i];
1088 1
                val = val % self->factors[i];
1089
            } else {
1090 1
                self->coordinates[i] = 0;
1091
            }
1092
        }
1093
    }
1094 1
    return PyArray_IntTupleFromIntp(nd, self->coordinates);
1095
}
1096

1097
static PyGetSetDef iter_getsets[] = {
1098
    {"coords",
1099
        (getter)iter_coords_get,
1100
        NULL,
1101
        NULL, NULL},
1102
    {NULL, NULL, NULL, NULL, NULL},
1103
};
1104

1105
NPY_NO_EXPORT PyTypeObject PyArrayIter_Type = {
1106
    PyVarObject_HEAD_INIT(NULL, 0)
1107
    .tp_name = "numpy.flatiter",
1108
    .tp_basicsize = sizeof(PyArrayIterObject),
1109
    .tp_dealloc = (destructor)arrayiter_dealloc,
1110
    .tp_as_mapping = &iter_as_mapping,
1111
    .tp_flags = Py_TPFLAGS_DEFAULT,
1112
    .tp_richcompare = (richcmpfunc)iter_richcompare,
1113
    .tp_iternext = (iternextfunc)arrayiter_next,
1114
    .tp_methods = iter_methods,
1115
    .tp_members = iter_members,
1116
    .tp_getset = iter_getsets,
1117
};
1118

1119
/** END of Array Iterator **/
1120

1121
/* Adjust dimensionality and strides for index object iterators
1122
   --- i.e. broadcast
1123
*/
1124
/*NUMPY_API*/
1125
NPY_NO_EXPORT int
1126 1
PyArray_Broadcast(PyArrayMultiIterObject *mit)
1127
{
1128
    int i, nd, k, j;
1129
    npy_intp tmp;
1130
    PyArrayIterObject *it;
1131

1132
    /* Discover the broadcast number of dimensions */
1133 1
    for (i = 0, nd = 0; i < mit->numiter; i++) {
1134 1
        nd = PyArray_MAX(nd, PyArray_NDIM(mit->iters[i]->ao));
1135
    }
1136 1
    mit->nd = nd;
1137

1138
    /* Discover the broadcast shape in each dimension */
1139 1
    for (i = 0; i < nd; i++) {
1140 1
        mit->dimensions[i] = 1;
1141 1
        for (j = 0; j < mit->numiter; j++) {
1142 1
            it = mit->iters[j];
1143
            /* This prepends 1 to shapes not already equal to nd */
1144 1
            k = i + PyArray_NDIM(it->ao) - nd;
1145 1
            if (k >= 0) {
1146 1
                tmp = PyArray_DIMS(it->ao)[k];
1147 1
                if (tmp == 1) {
1148 1
                    continue;
1149
                }
1150 1
                if (mit->dimensions[i] == 1) {
1151 1
                    mit->dimensions[i] = tmp;
1152
                }
1153 1
                else if (mit->dimensions[i] != tmp) {
1154 1
                    PyErr_SetString(PyExc_ValueError,
1155
                                    "shape mismatch: objects" \
1156
                                    " cannot be broadcast" \
1157
                                    " to a single shape");
1158 1
                    return -1;
1159
                }
1160
            }
1161
        }
1162
    }
1163

1164
    /*
1165
     * Reset the iterator dimensions and strides of each iterator
1166
     * object -- using 0 valued strides for broadcasting
1167
     * Need to check for overflow
1168
     */
1169 1
    tmp = PyArray_OverflowMultiplyList(mit->dimensions, mit->nd);
1170 1
    if (tmp < 0) {
1171 0
        PyErr_SetString(PyExc_ValueError,
1172
                        "broadcast dimensions too large.");
1173 0
        return -1;
1174
    }
1175 1
    mit->size = tmp;
1176 1
    for (i = 0; i < mit->numiter; i++) {
1177 1
        it = mit->iters[i];
1178 1
        it->nd_m1 = mit->nd - 1;
1179 1
        it->size = tmp;
1180 1
        nd = PyArray_NDIM(it->ao);
1181 1
        if (nd != 0) {
1182 1
            it->factors[mit->nd-1] = 1;
1183
        }
1184 1
        for (j = 0; j < mit->nd; j++) {
1185 1
            it->dims_m1[j] = mit->dimensions[j] - 1;
1186 1
            k = j + nd - mit->nd;
1187
            /*
1188
             * If this dimension was added or shape of
1189
             * underlying array was 1
1190
             */
1191 1
            if ((k < 0) ||
1192 1
                PyArray_DIMS(it->ao)[k] != mit->dimensions[j]) {
1193 1
                it->contiguous = 0;
1194 1
                it->strides[j] = 0;
1195
            }
1196
            else {
1197 1
                it->strides[j] = PyArray_STRIDES(it->ao)[k];
1198
            }
1199 1
            it->backstrides[j] = it->strides[j] * it->dims_m1[j];
1200 1
            if (j > 0)
1201 1
                it->factors[mit->nd-j-1] =
1202 1
                    it->factors[mit->nd-j] * mit->dimensions[mit->nd-j];
1203
        }
1204 1
        PyArray_ITER_RESET(it);
1205
    }
1206
    return 0;
1207
}
1208

1209
static NPY_INLINE PyObject*
1210
multiiter_wrong_number_of_args(void)
1211
{
1212 1
    return PyErr_Format(PyExc_ValueError,
1213
                        "Need at least 0 and at most %d "
1214
                        "array objects.", NPY_MAXARGS);
1215
}
1216

1217
/*
1218
 * Common implementation for all PyArrayMultiIterObject constructors.
1219
 */
1220
static PyObject*
1221 1
multiiter_new_impl(int n_args, PyObject **args)
1222
{
1223
    PyArrayMultiIterObject *multi;
1224
    int i;
1225

1226 1
    multi = PyArray_malloc(sizeof(PyArrayMultiIterObject));
1227 1
    if (multi == NULL) {
1228 0
        return PyErr_NoMemory();
1229
    }
1230 1
    PyObject_Init((PyObject *)multi, &PyArrayMultiIter_Type);
1231 1
    multi->numiter = 0;
1232

1233 1
    for (i = 0; i < n_args; ++i) {
1234 1
        PyObject *obj = args[i];
1235
        PyObject *arr;
1236
        PyArrayIterObject *it;
1237

1238 1
        if (PyObject_IsInstance(obj, (PyObject *)&PyArrayMultiIter_Type)) {
1239 1
            PyArrayMultiIterObject *mit = (PyArrayMultiIterObject *)obj;
1240
            int j;
1241

1242 1
            if (multi->numiter + mit->numiter > NPY_MAXARGS) {
1243
                multiiter_wrong_number_of_args();
1244
                goto fail;
1245
            }
1246 1
            for (j = 0; j < mit->numiter; ++j) {
1247 1
                arr = (PyObject *)mit->iters[j]->ao;
1248 1
                it = (PyArrayIterObject *)PyArray_IterNew(arr);
1249 1
                if (it == NULL) {
1250
                    goto fail;
1251
                }
1252 1
                multi->iters[multi->numiter++] = it;
1253
            }
1254
        }
1255 1
        else if (multi->numiter < NPY_MAXARGS) {
1256 1
            arr = PyArray_FromAny(obj, NULL, 0, 0, 0, NULL);
1257 1
            if (arr == NULL) {
1258
                goto fail;
1259
            }
1260 1
            it = (PyArrayIterObject *)PyArray_IterNew(arr);
1261 1
            Py_DECREF(arr);
1262 1
            if (it == NULL) {
1263
                goto fail;
1264
            }
1265 1
            multi->iters[multi->numiter++] = it;
1266
        }
1267
        else {
1268
            multiiter_wrong_number_of_args();
1269
            goto fail;
1270
        }
1271
    }
1272

1273 1
    if (multi->numiter < 0) {
1274
        multiiter_wrong_number_of_args();
1275
        goto fail;
1276
    }
1277 1
    if (PyArray_Broadcast(multi) < 0) {
1278
        goto fail;
1279
    }
1280 1
    PyArray_MultiIter_RESET(multi);
1281

1282
    return (PyObject *)multi;
1283

1284 1
fail:
1285 1
    Py_DECREF(multi);
1286

1287
    return NULL;
1288
}
1289

1290
/*NUMPY_API
1291
 * Get MultiIterator from array of Python objects and any additional
1292
 *
1293
 * PyObject **mps - array of PyObjects
1294
 * int n - number of PyObjects in the array
1295
 * int nadd - number of additional arrays to include in the iterator.
1296
 *
1297
 * Returns a multi-iterator object.
1298
 */
1299
NPY_NO_EXPORT PyObject*
1300 1
PyArray_MultiIterFromObjects(PyObject **mps, int n, int nadd, ...)
1301
{
1302
    PyObject *args_impl[NPY_MAXARGS];
1303 1
    int ntot = n + nadd;
1304
    int i;
1305
    va_list va;
1306

1307 1
    if ((ntot > NPY_MAXARGS) || (ntot < 0)) {
1308 0
        return multiiter_wrong_number_of_args();
1309
    }
1310

1311 1
    for (i = 0; i < n; ++i) {
1312 1
        args_impl[i] = mps[i];
1313
    }
1314

1315 1
    va_start(va, nadd);
1316 1
    for (; i < ntot; ++i) {
1317 1
        args_impl[i] = va_arg(va, PyObject *);
1318
    }
1319 1
    va_end(va);
1320

1321 1
    return multiiter_new_impl(ntot, args_impl);
1322
}
1323

1324
/*NUMPY_API
1325
 * Get MultiIterator,
1326
 */
1327
NPY_NO_EXPORT PyObject*
1328 1
PyArray_MultiIterNew(int n, ...)
1329
{
1330
    PyObject *args_impl[NPY_MAXARGS];
1331
    int i;
1332
    va_list va;
1333

1334 1
    if ((n > NPY_MAXARGS) || (n < 0)) {
1335 0
        return multiiter_wrong_number_of_args();
1336
    }
1337

1338 1
    va_start(va, n);
1339 1
    for (i = 0; i < n; ++i) {
1340 1
        args_impl[i] = va_arg(va, PyObject *);
1341
    }
1342 1
    va_end(va);
1343

1344 1
    return multiiter_new_impl(n, args_impl);
1345
}
1346

1347
static PyObject*
1348 1
arraymultiter_new(PyTypeObject *NPY_UNUSED(subtype), PyObject *args,
1349
                  PyObject *kwds)
1350
{
1351
    PyObject *ret, *fast_seq;
1352
    Py_ssize_t n;
1353

1354 1
    if (kwds != NULL && PyDict_Size(kwds) > 0) {
1355 1
        PyErr_SetString(PyExc_ValueError,
1356
                        "keyword arguments not accepted.");
1357 1
        return NULL;
1358
    }
1359

1360 1
    fast_seq = PySequence_Fast(args, "");  // needed for pypy
1361 1
    if (fast_seq == NULL) {
1362
        return NULL;
1363
    }
1364 1
    n = PySequence_Fast_GET_SIZE(fast_seq);
1365 1
    if (n > NPY_MAXARGS) {
1366 1
        Py_DECREF(fast_seq);
1367 1
        return multiiter_wrong_number_of_args();
1368
    }
1369 1
    ret = multiiter_new_impl(n, PySequence_Fast_ITEMS(fast_seq));
1370 1
    Py_DECREF(fast_seq);
1371
    return ret;
1372
}
1373

1374
static PyObject *
1375 1
arraymultiter_next(PyArrayMultiIterObject *multi)
1376
{
1377
    PyObject *ret;
1378
    int i, n;
1379

1380 1
    n = multi->numiter;
1381 1
    ret = PyTuple_New(n);
1382 1
    if (ret == NULL) {
1383
        return NULL;
1384
    }
1385 1
    if (multi->index < multi->size) {
1386 1
        for (i = 0; i < n; i++) {
1387 1
            PyArrayIterObject *it=multi->iters[i];
1388 1
            PyTuple_SET_ITEM(ret, i,
1389
                             PyArray_ToScalar(it->dataptr, it->ao));
1390 1
            PyArray_ITER_NEXT(it);
1391
        }
1392 1
        multi->index++;
1393 1
        return ret;
1394
    }
1395 1
    Py_DECREF(ret);
1396
    return NULL;
1397
}
1398

1399
static void
1400 1
arraymultiter_dealloc(PyArrayMultiIterObject *multi)
1401
{
1402
    int i;
1403

1404 1
    for (i = 0; i < multi->numiter; i++) {
1405 1
        Py_XDECREF(multi->iters[i]);
1406
    }
1407 1
    Py_TYPE(multi)->tp_free((PyObject *)multi);
1408
}
1409

1410
static PyObject *
1411 1
arraymultiter_size_get(PyArrayMultiIterObject *self)
1412
{
1413
#if NPY_SIZEOF_INTP <= NPY_SIZEOF_LONG
1414 1
    return PyLong_FromLong((long) self->size);
1415
#else
1416
    if (self->size < NPY_MAX_LONG) {
1417
        return PyLong_FromLong((long) self->size);
1418
    }
1419
    else {
1420
        return PyLong_FromLongLong((npy_longlong) self->size);
1421
    }
1422
#endif
1423
}
1424

1425
static PyObject *
1426 0
arraymultiter_index_get(PyArrayMultiIterObject *self)
1427
{
1428
#if NPY_SIZEOF_INTP <= NPY_SIZEOF_LONG
1429 0
    return PyLong_FromLong((long) self->index);
1430
#else
1431
    if (self->size < NPY_MAX_LONG) {
1432
        return PyLong_FromLong((long) self->index);
1433
    }
1434
    else {
1435
        return PyLong_FromLongLong((npy_longlong) self->index);
1436
    }
1437
#endif
1438
}
1439

1440
static PyObject *
1441 1
arraymultiter_shape_get(PyArrayMultiIterObject *self)
1442
{
1443 1
    return PyArray_IntTupleFromIntp(self->nd, self->dimensions);
1444
}
1445

1446
static PyObject *
1447 1
arraymultiter_iters_get(PyArrayMultiIterObject *self)
1448
{
1449
    PyObject *res;
1450
    int i, n;
1451

1452 1
    n = self->numiter;
1453 1
    res = PyTuple_New(n);
1454 1
    if (res == NULL) {
1455
        return res;
1456
    }
1457 1
    for (i = 0; i < n; i++) {
1458 1
        Py_INCREF(self->iters[i]);
1459 1
        PyTuple_SET_ITEM(res, i, (PyObject *)self->iters[i]);
1460
    }
1461
    return res;
1462
}
1463

1464
static PyGetSetDef arraymultiter_getsetlist[] = {
1465
    {"size",
1466
        (getter)arraymultiter_size_get,
1467
        NULL,
1468
        NULL, NULL},
1469
    {"index",
1470
        (getter)arraymultiter_index_get,
1471
        NULL,
1472
        NULL, NULL},
1473
    {"shape",
1474
        (getter)arraymultiter_shape_get,
1475
        NULL,
1476
        NULL, NULL},
1477
    {"iters",
1478
        (getter)arraymultiter_iters_get,
1479
        NULL,
1480
        NULL, NULL},
1481
    {NULL, NULL, NULL, NULL, NULL},
1482
};
1483

1484
static PyMemberDef arraymultiter_members[] = {
1485
    {"numiter",
1486
        T_INT,
1487
        offsetof(PyArrayMultiIterObject, numiter),
1488
        READONLY, NULL},
1489
    {"nd",
1490
        T_INT,
1491
        offsetof(PyArrayMultiIterObject, nd),
1492
        READONLY, NULL},
1493
    {"ndim",
1494
        T_INT,
1495
        offsetof(PyArrayMultiIterObject, nd),
1496
        READONLY, NULL},
1497
    {NULL, 0, 0, 0, NULL},
1498
};
1499

1500
static PyObject *
1501 0
arraymultiter_reset(PyArrayMultiIterObject *self, PyObject *args)
1502
{
1503 0
    if (!PyArg_ParseTuple(args, "")) {
1504
        return NULL;
1505
    }
1506 0
    PyArray_MultiIter_RESET(self);
1507 0
    Py_RETURN_NONE;
1508
}
1509

1510
static PyMethodDef arraymultiter_methods[] = {
1511
    {"reset",
1512
        (PyCFunction) arraymultiter_reset,
1513
        METH_VARARGS, NULL},
1514
    {NULL, NULL, 0, NULL},      /* sentinel */
1515
};
1516

1517
NPY_NO_EXPORT PyTypeObject PyArrayMultiIter_Type = {
1518
    PyVarObject_HEAD_INIT(NULL, 0)
1519
    .tp_name = "numpy.broadcast",
1520
    .tp_basicsize = sizeof(PyArrayMultiIterObject),
1521
    .tp_dealloc = (destructor)arraymultiter_dealloc,
1522
    .tp_flags = Py_TPFLAGS_DEFAULT,
1523
    .tp_iternext = (iternextfunc)arraymultiter_next,
1524
    .tp_methods = arraymultiter_methods,
1525
    .tp_members = arraymultiter_members,
1526
    .tp_getset = arraymultiter_getsetlist,
1527
    .tp_new = arraymultiter_new,
1528
};
1529

1530
/*========================= Neighborhood iterator ======================*/
1531

1532
static void neighiter_dealloc(PyArrayNeighborhoodIterObject* iter);
1533

1534 1
static char* _set_constant(PyArrayNeighborhoodIterObject* iter,
1535
        PyArrayObject *fill)
1536
{
1537
    char *ret;
1538 1
    PyArrayIterObject *ar = iter->_internal_iter;
1539
    int storeflags, st;
1540

1541 1
    ret = PyDataMem_NEW(PyArray_DESCR(ar->ao)->elsize);
1542 1
    if (ret == NULL) {
1543 0
        PyErr_SetNone(PyExc_MemoryError);
1544
        return NULL;
1545
    }
1546

1547 1
    if (PyArray_ISOBJECT(ar->ao)) {
1548 1
        memcpy(ret, PyArray_DATA(fill), sizeof(PyObject*));
1549 1
        Py_INCREF(*(PyObject**)ret);
1550
    } else {
1551
        /* Non-object types */
1552

1553 1
        storeflags = PyArray_FLAGS(ar->ao);
1554 1
        PyArray_ENABLEFLAGS(ar->ao, NPY_ARRAY_BEHAVED);
1555 1
        st = PyArray_SETITEM(ar->ao, ret, (PyObject*)fill);
1556 1
        ((PyArrayObject_fields *)ar->ao)->flags = storeflags;
1557

1558 1
        if (st < 0) {
1559 0
            PyDataMem_FREE(ret);
1560
            return NULL;
1561
        }
1562
    }
1563

1564
    return ret;
1565
}
1566

1567
#define _INF_SET_PTR(c) \
1568
    bd = coordinates[c] + p->coordinates[c]; \
1569
    if (bd < p->limits[c][0] || bd > p->limits[c][1]) { \
1570
        return niter->constant; \
1571
    } \
1572
    _coordinates[c] = bd;
1573

1574
/* set the dataptr from its current coordinates */
1575
static char*
1576 1
get_ptr_constant(PyArrayIterObject* _iter, const npy_intp *coordinates)
1577
{
1578
    int i;
1579
    npy_intp bd, _coordinates[NPY_MAXDIMS];
1580 1
    PyArrayNeighborhoodIterObject *niter = (PyArrayNeighborhoodIterObject*)_iter;
1581 1
    PyArrayIterObject *p = niter->_internal_iter;
1582

1583 1
    for(i = 0; i < niter->nd; ++i) {
1584 1
        _INF_SET_PTR(i)
1585
    }
1586

1587 1
    return p->translate(p, _coordinates);
1588
}
1589
#undef _INF_SET_PTR
1590

1591
#define _NPY_IS_EVEN(x) ((x) % 2 == 0)
1592

1593
/* For an array x of dimension n, and given index i, returns j, 0 <= j < n
1594
 * such as x[i] = x[j], with x assumed to be mirrored. For example, for x =
1595
 * {1, 2, 3} (n = 3)
1596
 *
1597
 * index -5 -4 -3 -2 -1 0 1 2 3 4 5 6
1598
 * value  2  3  3  2  1 1 2 3 3 2 1 1
1599
 *
1600
 * _npy_pos_index_mirror(4, 3) will return 1, because x[4] = x[1]*/
1601
NPY_INLINE static npy_intp
1602
__npy_pos_remainder(npy_intp i, npy_intp n)
1603
{
1604
    npy_intp k, l, j;
1605

1606
    /* Mirror i such as it is guaranteed to be positive */
1607 1
    if (i < 0) {
1608 1
        i = - i - 1;
1609
    }
1610

1611
    /* compute k and l such as i = k * n + l, 0 <= l < k */
1612 1
    k = i / n;
1613 1
    l = i - k * n;
1614

1615 1
    if (_NPY_IS_EVEN(k)) {
1616
        j = l;
1617
    } else {
1618 1
        j = n - 1 - l;
1619
    }
1620
    return j;
1621
}
1622
#undef _NPY_IS_EVEN
1623

1624
#define _INF_SET_PTR_MIRROR(c) \
1625
    lb = p->limits[c][0]; \
1626
    bd = coordinates[c] + p->coordinates[c] - lb; \
1627
    _coordinates[c] = lb + __npy_pos_remainder(bd, p->limits_sizes[c]);
1628

1629
/* set the dataptr from its current coordinates */
1630
static char*
1631 1
get_ptr_mirror(PyArrayIterObject* _iter, const npy_intp *coordinates)
1632
{
1633
    int i;
1634
    npy_intp bd, _coordinates[NPY_MAXDIMS], lb;
1635 1
    PyArrayNeighborhoodIterObject *niter = (PyArrayNeighborhoodIterObject*)_iter;
1636 1
    PyArrayIterObject *p = niter->_internal_iter;
1637

1638 1
    for(i = 0; i < niter->nd; ++i) {
1639 1
        _INF_SET_PTR_MIRROR(i)
1640
    }
1641

1642 1
    return p->translate(p, _coordinates);
1643
}
1644
#undef _INF_SET_PTR_MIRROR
1645

1646
/* compute l such as i = k * n + l, 0 <= l < |k| */
1647
NPY_INLINE static npy_intp
1648
__npy_euclidean_division(npy_intp i, npy_intp n)
1649
{
1650
    npy_intp l;
1651

1652 1
    l = i % n;
1653 1
    if (l < 0) {
1654 1
        l += n;
1655
    }
1656
    return l;
1657
}
1658

1659
#define _INF_SET_PTR_CIRCULAR(c) \
1660
    lb = p->limits[c][0]; \
1661
    bd = coordinates[c] + p->coordinates[c] - lb; \
1662
    _coordinates[c] = lb + __npy_euclidean_division(bd, p->limits_sizes[c]);
1663

1664
static char*
1665 1
get_ptr_circular(PyArrayIterObject* _iter, const npy_intp *coordinates)
1666
{
1667
    int i;
1668
    npy_intp bd, _coordinates[NPY_MAXDIMS], lb;
1669 1
    PyArrayNeighborhoodIterObject *niter = (PyArrayNeighborhoodIterObject*)_iter;
1670 1
    PyArrayIterObject *p = niter->_internal_iter;
1671

1672 1
    for(i = 0; i < niter->nd; ++i) {
1673 1
        _INF_SET_PTR_CIRCULAR(i)
1674
    }
1675 1
    return p->translate(p, _coordinates);
1676
}
1677

1678
#undef _INF_SET_PTR_CIRCULAR
1679

1680
/*
1681
 * fill and x->ao should have equivalent types
1682
 */
1683
/*NUMPY_API
1684
 * A Neighborhood Iterator object.
1685
*/
1686
NPY_NO_EXPORT PyObject*
1687 1
PyArray_NeighborhoodIterNew(PyArrayIterObject *x, const npy_intp *bounds,
1688
                            int mode, PyArrayObject* fill)
1689
{
1690
    int i;
1691
    PyArrayNeighborhoodIterObject *ret;
1692

1693 1
    ret = PyArray_malloc(sizeof(*ret));
1694 1
    if (ret == NULL) {
1695
        return NULL;
1696
    }
1697 1
    PyObject_Init((PyObject *)ret, &PyArrayNeighborhoodIter_Type);
1698

1699 1
    Py_INCREF(x->ao);  /* PyArray_RawIterBaseInit steals a reference */
1700 1
    PyArray_RawIterBaseInit((PyArrayIterObject*)ret, x->ao);
1701 1
    Py_INCREF(x);
1702 1
    ret->_internal_iter = x;
1703

1704 1
    ret->nd = PyArray_NDIM(x->ao);
1705

1706 1
    for (i = 0; i < ret->nd; ++i) {
1707 1
        ret->dimensions[i] = PyArray_DIMS(x->ao)[i];
1708
    }
1709

1710
    /* Compute the neighborhood size and copy the shape */
1711 1
    ret->size = 1;
1712 1
    for (i = 0; i < ret->nd; ++i) {
1713 1
        ret->bounds[i][0] = bounds[2 * i];
1714 1
        ret->bounds[i][1] = bounds[2 * i + 1];
1715 1
        ret->size *= (ret->bounds[i][1] - ret->bounds[i][0]) + 1;
1716

1717
        /* limits keep track of valid ranges for the neighborhood: if a bound
1718
         * of the neighborhood is outside the array, then limits is the same as
1719
         * boundaries. On the contrary, if a bound is strictly inside the
1720
         * array, then limits correspond to the array range. For example, for
1721
         * an array [1, 2, 3], if bounds are [-1, 3], limits will be [-1, 3],
1722
         * but if bounds are [1, 2], then limits will be [0, 2].
1723
         *
1724
         * This is used by neighborhood iterators stacked on top of this one */
1725 1
        ret->limits[i][0] = ret->bounds[i][0] < 0 ? ret->bounds[i][0] : 0;
1726 1
        ret->limits[i][1] = ret->bounds[i][1] >= ret->dimensions[i] - 1 ?
1727 1
                            ret->bounds[i][1] :
1728
                            ret->dimensions[i] - 1;
1729 1
        ret->limits_sizes[i] = (ret->limits[i][1] - ret->limits[i][0]) + 1;
1730
    }
1731

1732 1
    switch (mode) {
1733 1
        case NPY_NEIGHBORHOOD_ITER_ZERO_PADDING:
1734 1
            ret->constant = PyArray_Zero(x->ao);
1735 1
            ret->mode = mode;
1736 1
            ret->translate = &get_ptr_constant;
1737 1
            break;
1738 1
        case NPY_NEIGHBORHOOD_ITER_ONE_PADDING:
1739 1
            ret->constant = PyArray_One(x->ao);
1740 1
            ret->mode = mode;
1741 1
            ret->translate = &get_ptr_constant;
1742 1
            break;
1743 1
        case NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING:
1744
            /* New reference in returned value of _set_constant if array
1745
             * object */
1746
            assert(PyArray_EquivArrTypes(x->ao, fill) == NPY_TRUE);
1747 1
            ret->constant = _set_constant(ret, fill);
1748 1
            if (ret->constant == NULL) {
1749
                goto clean_x;
1750
            }
1751 1
            ret->mode = mode;
1752 1
            ret->translate = &get_ptr_constant;
1753 1
            break;
1754 1
        case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
1755 1
            ret->mode = mode;
1756 1
            ret->constant = NULL;
1757 1
            ret->translate = &get_ptr_mirror;
1758 1
            break;
1759 1
        case NPY_NEIGHBORHOOD_ITER_CIRCULAR_PADDING:
1760 1
            ret->mode = mode;
1761 1
            ret->constant = NULL;
1762 1
            ret->translate = &get_ptr_circular;
1763 1
            break;
1764 0
        default:
1765 0
            PyErr_SetString(PyExc_ValueError, "Unsupported padding mode");
1766 0
            goto clean_x;
1767
    }
1768

1769
    /*
1770
     * XXX: we force x iterator to be non contiguous because we need
1771
     * coordinates... Modifying the iterator here is not great
1772
     */
1773 1
    x->contiguous = 0;
1774

1775 1
    PyArrayNeighborhoodIter_Reset(ret);
1776

1777 1
    return (PyObject*)ret;
1778

1779 0
clean_x:
1780 0
    Py_DECREF(ret->_internal_iter);
1781 0
    array_iter_base_dealloc((PyArrayIterObject*)ret);
1782 0
    PyArray_free((PyArrayObject*)ret);
1783 0
    return NULL;
1784
}
1785

1786 1
static void neighiter_dealloc(PyArrayNeighborhoodIterObject* iter)
1787
{
1788 1
    if (iter->mode == NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING) {
1789 1
        if (PyArray_ISOBJECT(iter->_internal_iter->ao)) {
1790 1
            Py_DECREF(*(PyObject**)iter->constant);
1791
        }
1792
    }
1793 1
    PyDataMem_FREE(iter->constant);
1794 1
    Py_DECREF(iter->_internal_iter);
1795

1796 1
    array_iter_base_dealloc((PyArrayIterObject*)iter);
1797 1
    PyArray_free((PyArrayObject*)iter);
1798
}
1799

1800
NPY_NO_EXPORT PyTypeObject PyArrayNeighborhoodIter_Type = {
1801
    PyVarObject_HEAD_INIT(NULL, 0)
1802
    .tp_name = "numpy.neigh_internal_iter",
1803
    .tp_basicsize = sizeof(PyArrayNeighborhoodIterObject),
1804
    .tp_dealloc = (destructor)neighiter_dealloc,
1805
    .tp_flags = Py_TPFLAGS_DEFAULT,
1806
};

Read our documentation on viewing source code .

Loading