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
coerce_index(PyObject * o,npy_intp * v)30 coerce_index(PyObject *o, npy_intp *v)
31 {
32 *v = PyArray_PyIntAsIntp(o);
33
34 if ((*v) == -1 && PyErr_Occurred()) {
35 PyErr_Clear();
36 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
parse_index_entry(PyObject * op,npy_intp * step_size,npy_intp * n_steps,npy_intp max,int axis,int check_index)48 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 if (op == Py_None) {
55 *n_steps = NEWAXIS_INDEX;
56 i = 0;
57 }
58 else if (op == Py_Ellipsis) {
59 *n_steps = ELLIPSIS_INDEX;
60 i = 0;
61 }
62 else if (PySlice_Check(op)) {
63 npy_intp stop;
64 if (PySlice_GetIndicesEx(op, max, &i, &stop, step_size, n_steps) < 0) {
65 goto fail;
66 }
67 if (*n_steps <= 0) {
68 *n_steps = 0;
69 *step_size = 1;
70 i = 0;
71 }
72 }
73 else if (coerce_index(op, &i)) {
74 *n_steps = SINGLE_INDEX;
75 *step_size = 0;
76 if (check_index) {
77 if (check_and_adjust_index(&i, max, axis, NULL) < 0) {
78 goto fail;
79 }
80 }
81 }
82 else {
83 PyErr_SetString(PyExc_IndexError,
84 "each index entry must be either a "
85 "slice, an integer, Ellipsis, or "
86 "newaxis");
87 goto fail;
88 }
89 return i;
90
91 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*
get_ptr_simple(PyArrayIterObject * iter,const npy_intp * coordinates)102 get_ptr_simple(PyArrayIterObject* iter, const npy_intp *coordinates)
103 {
104 npy_intp i;
105 char *ret;
106
107 ret = PyArray_DATA(iter->ao);
108
109 for(i = 0; i < PyArray_NDIM(iter->ao); ++i) {
110 ret += coordinates[i] * iter->strides[i];
111 }
112
113 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
PyArray_RawIterBaseInit(PyArrayIterObject * it,PyArrayObject * ao)125 PyArray_RawIterBaseInit(PyArrayIterObject *it, PyArrayObject *ao)
126 {
127 int nd, i;
128
129 nd = PyArray_NDIM(ao);
130 PyArray_UpdateFlags(ao, NPY_ARRAY_C_CONTIGUOUS);
131 if (PyArray_ISCONTIGUOUS(ao)) {
132 it->contiguous = 1;
133 }
134 else {
135 it->contiguous = 0;
136 }
137 it->ao = ao;
138 it->size = PyArray_SIZE(ao);
139 it->nd_m1 = nd - 1;
140 if (nd != 0) {
141 it->factors[nd-1] = 1;
142 }
143 for (i = 0; i < nd; i++) {
144 it->dims_m1[i] = PyArray_DIMS(ao)[i] - 1;
145 it->strides[i] = PyArray_STRIDES(ao)[i];
146 it->backstrides[i] = it->strides[i] * it->dims_m1[i];
147 if (i > 0) {
148 it->factors[nd-i-1] = it->factors[nd-i] * PyArray_DIMS(ao)[nd-i];
149 }
150 it->bounds[i][0] = 0;
151 it->bounds[i][1] = PyArray_DIMS(ao)[i] - 1;
152 it->limits[i][0] = 0;
153 it->limits[i][1] = PyArray_DIMS(ao)[i] - 1;
154 it->limits_sizes[i] = it->limits[i][1] - it->limits[i][0] + 1;
155 }
156
157 it->translate = &get_ptr_simple;
158 PyArray_ITER_RESET(it);
159
160 return;
161 }
162
163 static void
array_iter_base_dealloc(PyArrayIterObject * it)164 array_iter_base_dealloc(PyArrayIterObject *it)
165 {
166 Py_XDECREF(it->ao);
167 }
168
169 /*NUMPY_API
170 * Get Iterator.
171 */
172 NPY_NO_EXPORT PyObject *
PyArray_IterNew(PyObject * obj)173 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 if (!PyArray_Check(obj)) {
183 PyErr_BadInternalCall();
184 return NULL;
185 }
186 ao = (PyArrayObject *)obj;
187
188 it = (PyArrayIterObject *)PyArray_malloc(sizeof(PyArrayIterObject));
189 PyObject_Init((PyObject *)it, &PyArrayIter_Type);
190 /* it = PyObject_New(PyArrayIterObject, &PyArrayIter_Type);*/
191 if (it == NULL) {
192 return NULL;
193 }
194
195 Py_INCREF(ao); /* PyArray_RawIterBaseInit steals a reference */
196 PyArray_RawIterBaseInit(it, ao);
197 return (PyObject *)it;
198 }
199
200 /*NUMPY_API
201 * Get Iterator broadcast to a particular shape
202 */
203 NPY_NO_EXPORT PyObject *
PyArray_BroadcastToShape(PyObject * obj,npy_intp * dims,int nd)204 PyArray_BroadcastToShape(PyObject *obj, npy_intp *dims, int nd)
205 {
206 PyArrayIterObject *it;
207 int i, diff, j, compat, k;
208 PyArrayObject *ao = (PyArrayObject *)obj;
209
210 if (PyArray_NDIM(ao) > nd) {
211 goto err;
212 }
213 compat = 1;
214 diff = j = nd - PyArray_NDIM(ao);
215 for (i = 0; i < PyArray_NDIM(ao); i++, j++) {
216 if (PyArray_DIMS(ao)[i] == 1) {
217 continue;
218 }
219 if (PyArray_DIMS(ao)[i] != dims[j]) {
220 compat = 0;
221 break;
222 }
223 }
224 if (!compat) {
225 goto err;
226 }
227 it = (PyArrayIterObject *)PyArray_malloc(sizeof(PyArrayIterObject));
228 if (it == NULL) {
229 return NULL;
230 }
231 PyObject_Init((PyObject *)it, &PyArrayIter_Type);
232
233 PyArray_UpdateFlags(ao, NPY_ARRAY_C_CONTIGUOUS);
234 if (PyArray_ISCONTIGUOUS(ao)) {
235 it->contiguous = 1;
236 }
237 else {
238 it->contiguous = 0;
239 }
240 Py_INCREF(ao);
241 it->ao = ao;
242 it->size = PyArray_MultiplyList(dims, nd);
243 it->nd_m1 = nd - 1;
244 if (nd != 0) {
245 it->factors[nd-1] = 1;
246 }
247 for (i = 0; i < nd; i++) {
248 it->dims_m1[i] = dims[i] - 1;
249 k = i - diff;
250 if ((k < 0) || PyArray_DIMS(ao)[k] != dims[i]) {
251 it->contiguous = 0;
252 it->strides[i] = 0;
253 }
254 else {
255 it->strides[i] = PyArray_STRIDES(ao)[k];
256 }
257 it->backstrides[i] = it->strides[i] * it->dims_m1[i];
258 if (i > 0) {
259 it->factors[nd-i-1] = it->factors[nd-i] * dims[nd-i];
260 }
261 }
262 PyArray_ITER_RESET(it);
263 return (PyObject *)it;
264
265 err:
266 PyErr_SetString(PyExc_ValueError, "array is not broadcastable to "\
267 "correct shape");
268 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 *
PyArray_IterAllButAxis(PyObject * obj,int * inaxis)281 PyArray_IterAllButAxis(PyObject *obj, int *inaxis)
282 {
283 PyArrayObject *arr;
284 PyArrayIterObject *it;
285 int axis;
286
287 if (!PyArray_Check(obj)) {
288 PyErr_SetString(PyExc_ValueError,
289 "Numpy IterAllButAxis requires an ndarray");
290 return NULL;
291 }
292 arr = (PyArrayObject *)obj;
293
294 it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)arr);
295 if (it == NULL) {
296 return NULL;
297 }
298 if (PyArray_NDIM(arr)==0) {
299 return (PyObject *)it;
300 }
301 if (*inaxis < 0) {
302 int i, minaxis = 0;
303 npy_intp minstride = 0;
304 i = 0;
305 while (minstride == 0 && i < PyArray_NDIM(arr)) {
306 minstride = PyArray_STRIDE(arr,i);
307 i++;
308 }
309 for (i = 1; i < PyArray_NDIM(arr); i++) {
310 if (PyArray_STRIDE(arr,i) > 0 &&
311 PyArray_STRIDE(arr, i) < minstride) {
312 minaxis = i;
313 minstride = PyArray_STRIDE(arr,i);
314 }
315 }
316 *inaxis = minaxis;
317 }
318 axis = *inaxis;
319 /* adjust so that will not iterate over axis */
320 it->contiguous = 0;
321 if (it->size != 0) {
322 it->size /= PyArray_DIM(arr,axis);
323 }
324 it->dims_m1[axis] = 0;
325 it->backstrides[axis] = 0;
326
327 /*
328 * (won't fix factors so don't use
329 * PyArray_ITER_GOTO1D with this iterator)
330 */
331 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
PyArray_RemoveSmallest(PyArrayMultiIterObject * multi)343 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 if (multi->nd == 0) {
352 return -1;
353 }
354 for (i = 0; i < multi->nd; i++) {
355 sumstrides[i] = 0;
356 for (j = 0; j < multi->numiter; j++) {
357 sumstrides[i] += multi->iters[j]->strides[i];
358 }
359 }
360 axis = 0;
361 smallest = sumstrides[0];
362 /* Find longest dimension */
363 for (i = 1; i < multi->nd; i++) {
364 if (sumstrides[i] < smallest) {
365 axis = i;
366 smallest = sumstrides[i];
367 }
368 }
369 for(i = 0; i < multi->numiter; i++) {
370 it = multi->iters[i];
371 it->contiguous = 0;
372 if (it->size != 0) {
373 it->size /= (it->dims_m1[axis]+1);
374 }
375 it->dims_m1[axis] = 0;
376 it->backstrides[axis] = 0;
377 }
378 multi->size = multi->iters[0]->size;
379 return axis;
380 }
381
382 /* Returns an array scalar holding the element desired */
383
384 static PyObject *
arrayiter_next(PyArrayIterObject * it)385 arrayiter_next(PyArrayIterObject *it)
386 {
387 PyObject *ret;
388
389 if (it->index < it->size) {
390 ret = PyArray_ToScalar(it->dataptr, it->ao);
391 PyArray_ITER_NEXT(it);
392 return ret;
393 }
394 return NULL;
395 }
396
397 static void
arrayiter_dealloc(PyArrayIterObject * it)398 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 array_iter_base_dealloc(it);
405 PyArray_free(it);
406 }
407
408 static Py_ssize_t
iter_length(PyArrayIterObject * self)409 iter_length(PyArrayIterObject *self)
410 {
411 return self->size;
412 }
413
414
415 static PyArrayObject *
iter_subscript_Bool(PyArrayIterObject * self,PyArrayObject * ind)416 iter_subscript_Bool(PyArrayIterObject *self, PyArrayObject *ind)
417 {
418 npy_intp counter, strides;
419 int itemsize;
420 npy_intp count = 0;
421 char *dptr, *optr;
422 PyArrayObject *ret;
423 int swap;
424 PyArray_CopySwapFunc *copyswap;
425
426
427 if (PyArray_NDIM(ind) != 1) {
428 PyErr_SetString(PyExc_ValueError,
429 "boolean index array should have 1 dimension");
430 return NULL;
431 }
432 counter = PyArray_DIMS(ind)[0];
433 if (counter > self->size) {
434 PyErr_SetString(PyExc_ValueError,
435 "too many boolean indices");
436 return NULL;
437 }
438
439 strides = PyArray_STRIDES(ind)[0];
440 dptr = PyArray_DATA(ind);
441 /* Get size of return array */
442 while (counter--) {
443 if (*((npy_bool *)dptr) != 0) {
444 count++;
445 }
446 dptr += strides;
447 }
448 itemsize = PyArray_DESCR(self->ao)->elsize;
449 Py_INCREF(PyArray_DESCR(self->ao));
450 ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self->ao),
451 PyArray_DESCR(self->ao), 1, &count,
452 NULL, NULL,
453 0, (PyObject *)self->ao);
454 if (ret == NULL) {
455 return NULL;
456 }
457 /* Set up loop */
458 optr = PyArray_DATA(ret);
459 counter = PyArray_DIMS(ind)[0];
460 dptr = PyArray_DATA(ind);
461 copyswap = PyArray_DESCR(self->ao)->f->copyswap;
462 /* Loop over Boolean array */
463 swap = (PyArray_ISNOTSWAPPED(self->ao) != PyArray_ISNOTSWAPPED(ret));
464 while (counter--) {
465 if (*((npy_bool *)dptr) != 0) {
466 copyswap(optr, self->dataptr, swap, self->ao);
467 optr += itemsize;
468 }
469 dptr += strides;
470 PyArray_ITER_NEXT(self);
471 }
472 PyArray_ITER_RESET(self);
473 return ret;
474 }
475
476 static PyObject *
iter_subscript_int(PyArrayIterObject * self,PyArrayObject * ind)477 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 itemsize = PyArray_DESCR(self->ao)->elsize;
489 if (PyArray_NDIM(ind) == 0) {
490 num = *((npy_intp *)PyArray_DATA(ind));
491 if (check_and_adjust_index(&num, self->size, -1, NULL) < 0) {
492 PyArray_ITER_RESET(self);
493 return NULL;
494 }
495 else {
496 PyObject *tmp;
497 PyArray_ITER_GOTO1D(self, num);
498 tmp = PyArray_ToScalar(self->dataptr, self->ao);
499 PyArray_ITER_RESET(self);
500 return tmp;
501 }
502 }
503
504 Py_INCREF(PyArray_DESCR(self->ao));
505 ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self->ao),
506 PyArray_DESCR(self->ao),
507 PyArray_NDIM(ind),
508 PyArray_DIMS(ind),
509 NULL, NULL,
510 0, (PyObject *)self->ao);
511 if (ret == NULL) {
512 return NULL;
513 }
514 optr = PyArray_DATA(ret);
515 ind_it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)ind);
516 if (ind_it == NULL) {
517 Py_DECREF(ret);
518 return NULL;
519 }
520 counter = ind_it->size;
521 copyswap = PyArray_DESCR(ret)->f->copyswap;
522 swap = (PyArray_ISNOTSWAPPED(ret) != PyArray_ISNOTSWAPPED(self->ao));
523 while (counter--) {
524 num = *((npy_intp *)(ind_it->dataptr));
525 if (check_and_adjust_index(&num, self->size, -1, NULL) < 0) {
526 Py_DECREF(ind_it);
527 Py_DECREF(ret);
528 PyArray_ITER_RESET(self);
529 return NULL;
530 }
531 PyArray_ITER_GOTO1D(self, num);
532 copyswap(optr, self->dataptr, swap, ret);
533 optr += itemsize;
534 PyArray_ITER_NEXT(ind_it);
535 }
536 Py_DECREF(ind_it);
537 PyArray_ITER_RESET(self);
538 return (PyObject *)ret;
539 }
540
541 /* Always returns arrays */
542 NPY_NO_EXPORT PyObject *
iter_subscript(PyArrayIterObject * self,PyObject * ind)543 iter_subscript(PyArrayIterObject *self, PyObject *ind)
544 {
545 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 PyObject *obj = NULL;
553 PyObject *new;
554 PyArray_CopySwapFunc *copyswap;
555
556 if (ind == Py_Ellipsis) {
557 ind = PySlice_New(NULL, NULL, NULL);
558 obj = iter_subscript(self, ind);
559 Py_DECREF(ind);
560 return obj;
561 }
562 if (PyTuple_Check(ind)) {
563 int len;
564 len = PyTuple_GET_SIZE(ind);
565 if (len > 1) {
566 goto fail;
567 }
568 if (len == 0) {
569 Py_INCREF(self->ao);
570 return (PyObject *)self->ao;
571 }
572 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 PyArray_ITER_RESET(self);
581
582 if (PyBool_Check(ind)) {
583 if (PyObject_IsTrue(ind)) {
584 return PyArray_ToScalar(self->dataptr, self->ao);
585 }
586 else { /* empty array */
587 npy_intp ii = 0;
588 dtype = PyArray_DESCR(self->ao);
589 Py_INCREF(dtype);
590 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 if (PyLong_Check(ind) || PySlice_Check(ind)) {
601 start = parse_index_entry(ind, &step_size, &n_steps,
602 self->size, 0, 1);
603 if (start == -1) {
604 goto fail;
605 }
606 if (n_steps == ELLIPSIS_INDEX || n_steps == NEWAXIS_INDEX) {
607 PyErr_SetString(PyExc_IndexError,
608 "cannot use Ellipsis or newaxes here");
609 goto fail;
610 }
611 PyArray_ITER_GOTO1D(self, start);
612 if (n_steps == SINGLE_INDEX) { /* Integer */
613 PyObject *tmp;
614 tmp = PyArray_ToScalar(self->dataptr, self->ao);
615 PyArray_ITER_RESET(self);
616 return tmp;
617 }
618 size = PyArray_DESCR(self->ao)->elsize;
619 dtype = PyArray_DESCR(self->ao);
620 Py_INCREF(dtype);
621 ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self->ao),
622 dtype,
623 1, &n_steps,
624 NULL, NULL,
625 0, (PyObject *)self->ao);
626 if (ret == NULL) {
627 goto fail;
628 }
629 dptr = PyArray_DATA(ret);
630 copyswap = PyArray_DESCR(ret)->f->copyswap;
631 while (n_steps--) {
632 copyswap(dptr, self->dataptr, 0, ret);
633 start += step_size;
634 PyArray_ITER_GOTO1D(self, start);
635 dptr += size;
636 }
637 PyArray_ITER_RESET(self);
638 return (PyObject *)ret;
639 }
640
641 /* convert to INTP array if Integer array scalar or List */
642 indtype = PyArray_DescrFromType(NPY_INTP);
643 if (PyArray_IsScalar(ind, Integer) || PyList_Check(ind)) {
644 Py_INCREF(indtype);
645 obj = PyArray_FromAny(ind, indtype, 0, 0, NPY_ARRAY_FORCECAST, NULL);
646 if (obj == NULL) {
647 goto fail;
648 }
649 }
650 else {
651 Py_INCREF(ind);
652 obj = ind;
653 }
654
655 /* Any remaining valid input is an array or has been turned into one */
656 if (!PyArray_Check(obj)) {
657 goto fail;
658 }
659
660 /* Check for Boolean array */
661 if (PyArray_TYPE((PyArrayObject *)obj) == NPY_BOOL) {
662 ret = iter_subscript_Bool(self, (PyArrayObject *)obj);
663 Py_DECREF(indtype);
664 Py_DECREF(obj);
665 return (PyObject *)ret;
666 }
667
668 /* Only integer arrays left */
669 if (!PyArray_ISINTEGER((PyArrayObject *)obj)) {
670 goto fail;
671 }
672
673 Py_INCREF(indtype);
674 new = PyArray_FromAny(obj, indtype, 0, 0,
675 NPY_ARRAY_FORCECAST | NPY_ARRAY_ALIGNED, NULL);
676 if (new == NULL) {
677 goto fail;
678 }
679 Py_DECREF(indtype);
680 Py_DECREF(obj);
681 ret = (PyArrayObject *)iter_subscript_int(self, (PyArrayObject *)new);
682 Py_DECREF(new);
683 return (PyObject *)ret;
684
685
686 fail:
687 if (!PyErr_Occurred()) {
688 PyErr_SetString(PyExc_IndexError, "unsupported iterator index");
689 }
690 Py_XDECREF(indtype);
691 Py_XDECREF(obj);
692 return NULL;
693
694 }
695
696
697 static int
iter_ass_sub_Bool(PyArrayIterObject * self,PyArrayObject * ind,PyArrayIterObject * val,int swap)698 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 if (PyArray_NDIM(ind) != 1) {
706 PyErr_SetString(PyExc_ValueError,
707 "boolean index array should have 1 dimension");
708 return -1;
709 }
710
711 counter = PyArray_DIMS(ind)[0];
712 if (counter > self->size) {
713 PyErr_SetString(PyExc_ValueError,
714 "boolean index array has too many values");
715 return -1;
716 }
717
718 strides = PyArray_STRIDES(ind)[0];
719 dptr = PyArray_DATA(ind);
720 PyArray_ITER_RESET(self);
721 /* Loop over Boolean array */
722 copyswap = PyArray_DESCR(self->ao)->f->copyswap;
723 while (counter--) {
724 if (*((npy_bool *)dptr) != 0) {
725 copyswap(self->dataptr, val->dataptr, swap, self->ao);
726 PyArray_ITER_NEXT(val);
727 if (val->index == val->size) {
728 PyArray_ITER_RESET(val);
729 }
730 }
731 dptr += strides;
732 PyArray_ITER_NEXT(self);
733 }
734 PyArray_ITER_RESET(self);
735 return 0;
736 }
737
738 static int
iter_ass_sub_int(PyArrayIterObject * self,PyArrayObject * ind,PyArrayIterObject * val,int swap)739 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 copyswap = PyArray_DESCR(self->ao)->f->copyswap;
748 if (PyArray_NDIM(ind) == 0) {
749 num = *((npy_intp *)PyArray_DATA(ind));
750 if (check_and_adjust_index(&num, self->size, -1, NULL) < 0) {
751 return -1;
752 }
753 PyArray_ITER_GOTO1D(self, num);
754 copyswap(self->dataptr, val->dataptr, swap, self->ao);
755 return 0;
756 }
757 ind_it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)ind);
758 if (ind_it == NULL) {
759 return -1;
760 }
761 counter = ind_it->size;
762 while (counter--) {
763 num = *((npy_intp *)(ind_it->dataptr));
764 if (check_and_adjust_index(&num, self->size, -1, NULL) < 0) {
765 Py_DECREF(ind_it);
766 return -1;
767 }
768 PyArray_ITER_GOTO1D(self, num);
769 copyswap(self->dataptr, val->dataptr, swap, self->ao);
770 PyArray_ITER_NEXT(ind_it);
771 PyArray_ITER_NEXT(val);
772 if (val->index == val->size) {
773 PyArray_ITER_RESET(val);
774 }
775 }
776 Py_DECREF(ind_it);
777 return 0;
778 }
779
780 NPY_NO_EXPORT int
iter_ass_subscript(PyArrayIterObject * self,PyObject * ind,PyObject * val)781 iter_ass_subscript(PyArrayIterObject *self, PyObject *ind, PyObject *val)
782 {
783 PyArrayObject *arrval = NULL;
784 PyArrayIterObject *val_it = NULL;
785 PyArray_Descr *type;
786 PyArray_Descr *indtype = NULL;
787 int swap, retval = -1;
788 npy_intp start, step_size;
789 npy_intp n_steps;
790 PyObject *obj = NULL;
791 PyArray_CopySwapFunc *copyswap;
792
793
794 if (val == NULL) {
795 PyErr_SetString(PyExc_TypeError,
796 "Cannot delete iterator elements");
797 return -1;
798 }
799
800 if (PyArray_FailUnlessWriteable(self->ao, "underlying array") < 0)
801 return -1;
802
803 if (ind == Py_Ellipsis) {
804 ind = PySlice_New(NULL, NULL, NULL);
805 retval = iter_ass_subscript(self, ind, val);
806 Py_DECREF(ind);
807 return retval;
808 }
809
810 if (PyTuple_Check(ind)) {
811 int len;
812 len = PyTuple_GET_SIZE(ind);
813 if (len > 1) {
814 goto finish;
815 }
816 ind = PyTuple_GET_ITEM(ind, 0);
817 }
818
819 type = PyArray_DESCR(self->ao);
820
821 /*
822 * Check for Boolean -- this is first because
823 * Bool is a subclass of Int
824 */
825 if (PyBool_Check(ind)) {
826 retval = 0;
827 if (PyObject_IsTrue(ind)) {
828 retval = PyArray_Pack(PyArray_DESCR(self->ao), self->dataptr, val);
829 }
830 goto finish;
831 }
832
833 if (PySequence_Check(ind) || PySlice_Check(ind)) {
834 goto skip;
835 }
836 start = PyArray_PyIntAsIntp(ind);
837 if (error_converting(start)) {
838 PyErr_Clear();
839 }
840 else {
841 if (check_and_adjust_index(&start, self->size, -1, NULL) < 0) {
842 goto finish;
843 }
844 PyArray_ITER_GOTO1D(self, start);
845 retval = PyArray_Pack(PyArray_DESCR(self->ao), self->dataptr, val);
846 PyArray_ITER_RESET(self);
847 if (retval < 0) {
848 PyErr_SetString(PyExc_ValueError,
849 "Error setting single item of array.");
850 }
851 goto finish;
852 }
853
854 skip:
855 Py_INCREF(type);
856 arrval = (PyArrayObject *)PyArray_FromAny(val, type, 0, 0,
857 NPY_ARRAY_FORCECAST, NULL);
858 if (arrval == NULL) {
859 return -1;
860 }
861 val_it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)arrval);
862 if (val_it == NULL) {
863 goto finish;
864 }
865 if (val_it->size == 0) {
866 retval = 0;
867 goto finish;
868 }
869
870 copyswap = PyArray_DESCR(arrval)->f->copyswap;
871 swap = (PyArray_ISNOTSWAPPED(self->ao)!=PyArray_ISNOTSWAPPED(arrval));
872
873 /* Check Slice */
874 if (PySlice_Check(ind)) {
875 start = parse_index_entry(ind, &step_size, &n_steps, self->size, 0, 0);
876 if (start == -1) {
877 goto finish;
878 }
879 if (n_steps == ELLIPSIS_INDEX || n_steps == NEWAXIS_INDEX) {
880 PyErr_SetString(PyExc_IndexError,
881 "cannot use Ellipsis or newaxes here");
882 goto finish;
883 }
884 PyArray_ITER_GOTO1D(self, start);
885 if (n_steps == SINGLE_INDEX) {
886 /* Integer */
887 copyswap(self->dataptr, PyArray_DATA(arrval), swap, arrval);
888 PyArray_ITER_RESET(self);
889 retval = 0;
890 goto finish;
891 }
892 while (n_steps--) {
893 copyswap(self->dataptr, val_it->dataptr, swap, arrval);
894 start += step_size;
895 PyArray_ITER_GOTO1D(self, start);
896 PyArray_ITER_NEXT(val_it);
897 if (val_it->index == val_it->size) {
898 PyArray_ITER_RESET(val_it);
899 }
900 }
901 PyArray_ITER_RESET(self);
902 retval = 0;
903 goto finish;
904 }
905
906 /* convert to INTP array if Integer array scalar or List */
907 indtype = PyArray_DescrFromType(NPY_INTP);
908 if (PyList_Check(ind)) {
909 Py_INCREF(indtype);
910 obj = PyArray_FromAny(ind, indtype, 0, 0, NPY_ARRAY_FORCECAST, NULL);
911 }
912 else {
913 Py_INCREF(ind);
914 obj = ind;
915 }
916
917 if (obj != NULL && PyArray_Check(obj)) {
918 /* Check for Boolean object */
919 if (PyArray_TYPE((PyArrayObject *)obj)==NPY_BOOL) {
920 if (iter_ass_sub_Bool(self, (PyArrayObject *)obj,
921 val_it, swap) < 0) {
922 goto finish;
923 }
924 retval=0;
925 }
926 /* Check for integer array */
927 else if (PyArray_ISINTEGER((PyArrayObject *)obj)) {
928 PyObject *new;
929 Py_INCREF(indtype);
930 new = PyArray_CheckFromAny(obj, indtype, 0, 0,
931 NPY_ARRAY_FORCECAST | NPY_ARRAY_BEHAVED_NS, NULL);
932 Py_DECREF(obj);
933 obj = new;
934 if (new == NULL) {
935 goto finish;
936 }
937 if (iter_ass_sub_int(self, (PyArrayObject *)obj,
938 val_it, swap) < 0) {
939 goto finish;
940 }
941 retval = 0;
942 }
943 }
944
945 finish:
946 if (!PyErr_Occurred() && retval < 0) {
947 PyErr_SetString(PyExc_IndexError, "unsupported iterator index");
948 }
949 Py_XDECREF(indtype);
950 Py_XDECREF(obj);
951 Py_XDECREF(val_it);
952 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 *
iter_array(PyArrayIterObject * it,PyObject * NPY_UNUSED (op))988 iter_array(PyArrayIterObject *it, PyObject *NPY_UNUSED(op))
989 {
990
991 PyArrayObject *ret;
992 npy_intp size;
993
994 size = PyArray_SIZE(it->ao);
995 Py_INCREF(PyArray_DESCR(it->ao));
996
997 if (PyArray_ISCONTIGUOUS(it->ao)) {
998 ret = (PyArrayObject *)PyArray_NewFromDescrAndBase(
999 &PyArray_Type, PyArray_DESCR(it->ao),
1000 1, &size, NULL, PyArray_DATA(it->ao),
1001 PyArray_FLAGS(it->ao), (PyObject *)it->ao, (PyObject *)it->ao);
1002 if (ret == NULL) {
1003 return NULL;
1004 }
1005 }
1006 else {
1007 ret = (PyArrayObject *)PyArray_NewFromDescr(
1008 &PyArray_Type, PyArray_DESCR(it->ao), 1, &size,
1009 NULL, NULL, 0,
1010 (PyObject *)it->ao);
1011 if (ret == NULL) {
1012 return NULL;
1013 }
1014 if (PyArray_CopyAnyInto(ret, it->ao) < 0) {
1015 Py_DECREF(ret);
1016 return NULL;
1017 }
1018 PyArray_CLEARFLAGS(ret, NPY_ARRAY_WRITEABLE);
1019 }
1020 return ret;
1021
1022 }
1023
1024 static PyObject *
iter_copy(PyArrayIterObject * it,PyObject * args)1025 iter_copy(PyArrayIterObject *it, PyObject *args)
1026 {
1027 if (!PyArg_ParseTuple(args, "")) {
1028 return NULL;
1029 }
1030 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 *
iter_richcompare(PyArrayIterObject * self,PyObject * other,int cmp_op)1045 iter_richcompare(PyArrayIterObject *self, PyObject *other, int cmp_op)
1046 {
1047 PyArrayObject *new;
1048 PyObject *ret;
1049 new = (PyArrayObject *)iter_array(self, NULL);
1050 if (new == NULL) {
1051 return NULL;
1052 }
1053 ret = array_richcompare(new, other, cmp_op);
1054 PyArray_ResolveWritebackIfCopy(new);
1055 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 *
iter_coords_get(PyArrayIterObject * self)1073 iter_coords_get(PyArrayIterObject *self)
1074 {
1075 int nd;
1076 nd = PyArray_NDIM(self->ao);
1077 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 val = self->index;
1085 for (i = 0; i < nd; i++) {
1086 if (self->factors[i] != 0) {
1087 self->coordinates[i] = val / self->factors[i];
1088 val = val % self->factors[i];
1089 } else {
1090 self->coordinates[i] = 0;
1091 }
1092 }
1093 }
1094 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
PyArray_Broadcast(PyArrayMultiIterObject * mit)1126 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 for (i = 0, nd = 0; i < mit->numiter; i++) {
1134 nd = PyArray_MAX(nd, PyArray_NDIM(mit->iters[i]->ao));
1135 }
1136 mit->nd = nd;
1137
1138 /* Discover the broadcast shape in each dimension */
1139 for (i = 0; i < nd; i++) {
1140 mit->dimensions[i] = 1;
1141 for (j = 0; j < mit->numiter; j++) {
1142 it = mit->iters[j];
1143 /* This prepends 1 to shapes not already equal to nd */
1144 k = i + PyArray_NDIM(it->ao) - nd;
1145 if (k >= 0) {
1146 tmp = PyArray_DIMS(it->ao)[k];
1147 if (tmp == 1) {
1148 continue;
1149 }
1150 if (mit->dimensions[i] == 1) {
1151 mit->dimensions[i] = tmp;
1152 }
1153 else if (mit->dimensions[i] != tmp) {
1154 PyErr_SetString(PyExc_ValueError,
1155 "shape mismatch: objects" \
1156 " cannot be broadcast" \
1157 " to a single shape");
1158 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 tmp = PyArray_OverflowMultiplyList(mit->dimensions, mit->nd);
1170 if (tmp < 0) {
1171 PyErr_SetString(PyExc_ValueError,
1172 "broadcast dimensions too large.");
1173 return -1;
1174 }
1175 mit->size = tmp;
1176 for (i = 0; i < mit->numiter; i++) {
1177 it = mit->iters[i];
1178 it->nd_m1 = mit->nd - 1;
1179 it->size = tmp;
1180 nd = PyArray_NDIM(it->ao);
1181 if (nd != 0) {
1182 it->factors[mit->nd-1] = 1;
1183 }
1184 for (j = 0; j < mit->nd; j++) {
1185 it->dims_m1[j] = mit->dimensions[j] - 1;
1186 k = j + nd - mit->nd;
1187 /*
1188 * If this dimension was added or shape of
1189 * underlying array was 1
1190 */
1191 if ((k < 0) ||
1192 PyArray_DIMS(it->ao)[k] != mit->dimensions[j]) {
1193 it->contiguous = 0;
1194 it->strides[j] = 0;
1195 }
1196 else {
1197 it->strides[j] = PyArray_STRIDES(it->ao)[k];
1198 }
1199 it->backstrides[j] = it->strides[j] * it->dims_m1[j];
1200 if (j > 0)
1201 it->factors[mit->nd-j-1] =
1202 it->factors[mit->nd-j] * mit->dimensions[mit->nd-j];
1203 }
1204 PyArray_ITER_RESET(it);
1205 }
1206 return 0;
1207 }
1208
1209 static NPY_INLINE PyObject*
multiiter_wrong_number_of_args(void)1210 multiiter_wrong_number_of_args(void)
1211 {
1212 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*
multiiter_new_impl(int n_args,PyObject ** args)1221 multiiter_new_impl(int n_args, PyObject **args)
1222 {
1223 PyArrayMultiIterObject *multi;
1224 int i;
1225
1226 multi = PyArray_malloc(sizeof(PyArrayMultiIterObject));
1227 if (multi == NULL) {
1228 return PyErr_NoMemory();
1229 }
1230 PyObject_Init((PyObject *)multi, &PyArrayMultiIter_Type);
1231 multi->numiter = 0;
1232
1233 for (i = 0; i < n_args; ++i) {
1234 PyObject *obj = args[i];
1235 PyObject *arr;
1236 PyArrayIterObject *it;
1237
1238 if (PyObject_IsInstance(obj, (PyObject *)&PyArrayMultiIter_Type)) {
1239 PyArrayMultiIterObject *mit = (PyArrayMultiIterObject *)obj;
1240 int j;
1241
1242 if (multi->numiter + mit->numiter > NPY_MAXARGS) {
1243 multiiter_wrong_number_of_args();
1244 goto fail;
1245 }
1246 for (j = 0; j < mit->numiter; ++j) {
1247 arr = (PyObject *)mit->iters[j]->ao;
1248 it = (PyArrayIterObject *)PyArray_IterNew(arr);
1249 if (it == NULL) {
1250 goto fail;
1251 }
1252 multi->iters[multi->numiter++] = it;
1253 }
1254 }
1255 else if (multi->numiter < NPY_MAXARGS) {
1256 arr = PyArray_FromAny(obj, NULL, 0, 0, 0, NULL);
1257 if (arr == NULL) {
1258 goto fail;
1259 }
1260 it = (PyArrayIterObject *)PyArray_IterNew(arr);
1261 Py_DECREF(arr);
1262 if (it == NULL) {
1263 goto fail;
1264 }
1265 multi->iters[multi->numiter++] = it;
1266 }
1267 else {
1268 multiiter_wrong_number_of_args();
1269 goto fail;
1270 }
1271 }
1272
1273 if (multi->numiter < 0) {
1274 multiiter_wrong_number_of_args();
1275 goto fail;
1276 }
1277 if (PyArray_Broadcast(multi) < 0) {
1278 goto fail;
1279 }
1280 PyArray_MultiIter_RESET(multi);
1281
1282 return (PyObject *)multi;
1283
1284 fail:
1285 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*
PyArray_MultiIterFromObjects(PyObject ** mps,int n,int nadd,...)1300 PyArray_MultiIterFromObjects(PyObject **mps, int n, int nadd, ...)
1301 {
1302 PyObject *args_impl[NPY_MAXARGS];
1303 int ntot = n + nadd;
1304 int i;
1305 va_list va;
1306
1307 if ((ntot > NPY_MAXARGS) || (ntot < 0)) {
1308 return multiiter_wrong_number_of_args();
1309 }
1310
1311 for (i = 0; i < n; ++i) {
1312 args_impl[i] = mps[i];
1313 }
1314
1315 va_start(va, nadd);
1316 for (; i < ntot; ++i) {
1317 args_impl[i] = va_arg(va, PyObject *);
1318 }
1319 va_end(va);
1320
1321 return multiiter_new_impl(ntot, args_impl);
1322 }
1323
1324 /*NUMPY_API
1325 * Get MultiIterator,
1326 */
1327 NPY_NO_EXPORT PyObject*
PyArray_MultiIterNew(int n,...)1328 PyArray_MultiIterNew(int n, ...)
1329 {
1330 PyObject *args_impl[NPY_MAXARGS];
1331 int i;
1332 va_list va;
1333
1334 if ((n > NPY_MAXARGS) || (n < 0)) {
1335 return multiiter_wrong_number_of_args();
1336 }
1337
1338 va_start(va, n);
1339 for (i = 0; i < n; ++i) {
1340 args_impl[i] = va_arg(va, PyObject *);
1341 }
1342 va_end(va);
1343
1344 return multiiter_new_impl(n, args_impl);
1345 }
1346
1347 static PyObject*
arraymultiter_new(PyTypeObject * NPY_UNUSED (subtype),PyObject * args,PyObject * kwds)1348 arraymultiter_new(PyTypeObject *NPY_UNUSED(subtype), PyObject *args,
1349 PyObject *kwds)
1350 {
1351 PyObject *ret, *fast_seq;
1352 Py_ssize_t n;
1353
1354 if (kwds != NULL && PyDict_Size(kwds) > 0) {
1355 PyErr_SetString(PyExc_ValueError,
1356 "keyword arguments not accepted.");
1357 return NULL;
1358 }
1359
1360 fast_seq = PySequence_Fast(args, ""); // needed for pypy
1361 if (fast_seq == NULL) {
1362 return NULL;
1363 }
1364 n = PySequence_Fast_GET_SIZE(fast_seq);
1365 if (n > NPY_MAXARGS) {
1366 Py_DECREF(fast_seq);
1367 return multiiter_wrong_number_of_args();
1368 }
1369 ret = multiiter_new_impl(n, PySequence_Fast_ITEMS(fast_seq));
1370 Py_DECREF(fast_seq);
1371 return ret;
1372 }
1373
1374 static PyObject *
arraymultiter_next(PyArrayMultiIterObject * multi)1375 arraymultiter_next(PyArrayMultiIterObject *multi)
1376 {
1377 PyObject *ret;
1378 int i, n;
1379
1380 n = multi->numiter;
1381 ret = PyTuple_New(n);
1382 if (ret == NULL) {
1383 return NULL;
1384 }
1385 if (multi->index < multi->size) {
1386 for (i = 0; i < n; i++) {
1387 PyArrayIterObject *it=multi->iters[i];
1388 PyTuple_SET_ITEM(ret, i,
1389 PyArray_ToScalar(it->dataptr, it->ao));
1390 PyArray_ITER_NEXT(it);
1391 }
1392 multi->index++;
1393 return ret;
1394 }
1395 Py_DECREF(ret);
1396 return NULL;
1397 }
1398
1399 static void
arraymultiter_dealloc(PyArrayMultiIterObject * multi)1400 arraymultiter_dealloc(PyArrayMultiIterObject *multi)
1401 {
1402 int i;
1403
1404 for (i = 0; i < multi->numiter; i++) {
1405 Py_XDECREF(multi->iters[i]);
1406 }
1407 Py_TYPE(multi)->tp_free((PyObject *)multi);
1408 }
1409
1410 static PyObject *
arraymultiter_size_get(PyArrayMultiIterObject * self)1411 arraymultiter_size_get(PyArrayMultiIterObject *self)
1412 {
1413 #if NPY_SIZEOF_INTP <= NPY_SIZEOF_LONG
1414 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 *
arraymultiter_index_get(PyArrayMultiIterObject * self)1426 arraymultiter_index_get(PyArrayMultiIterObject *self)
1427 {
1428 #if NPY_SIZEOF_INTP <= NPY_SIZEOF_LONG
1429 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 *
arraymultiter_shape_get(PyArrayMultiIterObject * self)1441 arraymultiter_shape_get(PyArrayMultiIterObject *self)
1442 {
1443 return PyArray_IntTupleFromIntp(self->nd, self->dimensions);
1444 }
1445
1446 static PyObject *
arraymultiter_iters_get(PyArrayMultiIterObject * self)1447 arraymultiter_iters_get(PyArrayMultiIterObject *self)
1448 {
1449 PyObject *res;
1450 int i, n;
1451
1452 n = self->numiter;
1453 res = PyTuple_New(n);
1454 if (res == NULL) {
1455 return res;
1456 }
1457 for (i = 0; i < n; i++) {
1458 Py_INCREF(self->iters[i]);
1459 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 *
arraymultiter_reset(PyArrayMultiIterObject * self,PyObject * args)1501 arraymultiter_reset(PyArrayMultiIterObject *self, PyObject *args)
1502 {
1503 if (!PyArg_ParseTuple(args, "")) {
1504 return NULL;
1505 }
1506 PyArray_MultiIter_RESET(self);
1507 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
_set_constant(PyArrayNeighborhoodIterObject * iter,PyArrayObject * fill)1534 static char* _set_constant(PyArrayNeighborhoodIterObject* iter,
1535 PyArrayObject *fill)
1536 {
1537 char *ret;
1538 PyArrayIterObject *ar = iter->_internal_iter;
1539 int storeflags, st;
1540
1541 ret = PyDataMem_NEW(PyArray_DESCR(ar->ao)->elsize);
1542 if (ret == NULL) {
1543 PyErr_SetNone(PyExc_MemoryError);
1544 return NULL;
1545 }
1546
1547 if (PyArray_ISOBJECT(ar->ao)) {
1548 memcpy(ret, PyArray_DATA(fill), sizeof(PyObject*));
1549 Py_INCREF(*(PyObject**)ret);
1550 } else {
1551 /* Non-object types */
1552
1553 storeflags = PyArray_FLAGS(ar->ao);
1554 PyArray_ENABLEFLAGS(ar->ao, NPY_ARRAY_BEHAVED);
1555 st = PyArray_SETITEM(ar->ao, ret, (PyObject*)fill);
1556 ((PyArrayObject_fields *)ar->ao)->flags = storeflags;
1557
1558 if (st < 0) {
1559 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*
get_ptr_constant(PyArrayIterObject * _iter,const npy_intp * coordinates)1576 get_ptr_constant(PyArrayIterObject* _iter, const npy_intp *coordinates)
1577 {
1578 int i;
1579 npy_intp bd, _coordinates[NPY_MAXDIMS];
1580 PyArrayNeighborhoodIterObject *niter = (PyArrayNeighborhoodIterObject*)_iter;
1581 PyArrayIterObject *p = niter->_internal_iter;
1582
1583 for(i = 0; i < niter->nd; ++i) {
1584 _INF_SET_PTR(i)
1585 }
1586
1587 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
__npy_pos_remainder(npy_intp i,npy_intp n)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 if (i < 0) {
1608 i = - i - 1;
1609 }
1610
1611 /* compute k and l such as i = k * n + l, 0 <= l < k */
1612 k = i / n;
1613 l = i - k * n;
1614
1615 if (_NPY_IS_EVEN(k)) {
1616 j = l;
1617 } else {
1618 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*
get_ptr_mirror(PyArrayIterObject * _iter,const npy_intp * coordinates)1631 get_ptr_mirror(PyArrayIterObject* _iter, const npy_intp *coordinates)
1632 {
1633 int i;
1634 npy_intp bd, _coordinates[NPY_MAXDIMS], lb;
1635 PyArrayNeighborhoodIterObject *niter = (PyArrayNeighborhoodIterObject*)_iter;
1636 PyArrayIterObject *p = niter->_internal_iter;
1637
1638 for(i = 0; i < niter->nd; ++i) {
1639 _INF_SET_PTR_MIRROR(i)
1640 }
1641
1642 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
__npy_euclidean_division(npy_intp i,npy_intp n)1648 __npy_euclidean_division(npy_intp i, npy_intp n)
1649 {
1650 npy_intp l;
1651
1652 l = i % n;
1653 if (l < 0) {
1654 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*
get_ptr_circular(PyArrayIterObject * _iter,const npy_intp * coordinates)1665 get_ptr_circular(PyArrayIterObject* _iter, const npy_intp *coordinates)
1666 {
1667 int i;
1668 npy_intp bd, _coordinates[NPY_MAXDIMS], lb;
1669 PyArrayNeighborhoodIterObject *niter = (PyArrayNeighborhoodIterObject*)_iter;
1670 PyArrayIterObject *p = niter->_internal_iter;
1671
1672 for(i = 0; i < niter->nd; ++i) {
1673 _INF_SET_PTR_CIRCULAR(i)
1674 }
1675 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*
PyArray_NeighborhoodIterNew(PyArrayIterObject * x,const npy_intp * bounds,int mode,PyArrayObject * fill)1687 PyArray_NeighborhoodIterNew(PyArrayIterObject *x, const npy_intp *bounds,
1688 int mode, PyArrayObject* fill)
1689 {
1690 int i;
1691 PyArrayNeighborhoodIterObject *ret;
1692
1693 ret = PyArray_malloc(sizeof(*ret));
1694 if (ret == NULL) {
1695 return NULL;
1696 }
1697 PyObject_Init((PyObject *)ret, &PyArrayNeighborhoodIter_Type);
1698
1699 Py_INCREF(x->ao); /* PyArray_RawIterBaseInit steals a reference */
1700 PyArray_RawIterBaseInit((PyArrayIterObject*)ret, x->ao);
1701 Py_INCREF(x);
1702 ret->_internal_iter = x;
1703
1704 ret->nd = PyArray_NDIM(x->ao);
1705
1706 for (i = 0; i < ret->nd; ++i) {
1707 ret->dimensions[i] = PyArray_DIMS(x->ao)[i];
1708 }
1709
1710 /* Compute the neighborhood size and copy the shape */
1711 ret->size = 1;
1712 for (i = 0; i < ret->nd; ++i) {
1713 ret->bounds[i][0] = bounds[2 * i];
1714 ret->bounds[i][1] = bounds[2 * i + 1];
1715 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 ret->limits[i][0] = ret->bounds[i][0] < 0 ? ret->bounds[i][0] : 0;
1726 ret->limits[i][1] = ret->bounds[i][1] >= ret->dimensions[i] - 1 ?
1727 ret->bounds[i][1] :
1728 ret->dimensions[i] - 1;
1729 ret->limits_sizes[i] = (ret->limits[i][1] - ret->limits[i][0]) + 1;
1730 }
1731
1732 switch (mode) {
1733 case NPY_NEIGHBORHOOD_ITER_ZERO_PADDING:
1734 ret->constant = PyArray_Zero(x->ao);
1735 ret->mode = mode;
1736 ret->translate = &get_ptr_constant;
1737 break;
1738 case NPY_NEIGHBORHOOD_ITER_ONE_PADDING:
1739 ret->constant = PyArray_One(x->ao);
1740 ret->mode = mode;
1741 ret->translate = &get_ptr_constant;
1742 break;
1743 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 ret->constant = _set_constant(ret, fill);
1748 if (ret->constant == NULL) {
1749 goto clean_x;
1750 }
1751 ret->mode = mode;
1752 ret->translate = &get_ptr_constant;
1753 break;
1754 case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
1755 ret->mode = mode;
1756 ret->constant = NULL;
1757 ret->translate = &get_ptr_mirror;
1758 break;
1759 case NPY_NEIGHBORHOOD_ITER_CIRCULAR_PADDING:
1760 ret->mode = mode;
1761 ret->constant = NULL;
1762 ret->translate = &get_ptr_circular;
1763 break;
1764 default:
1765 PyErr_SetString(PyExc_ValueError, "Unsupported padding mode");
1766 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 x->contiguous = 0;
1774
1775 PyArrayNeighborhoodIter_Reset(ret);
1776
1777 return (PyObject*)ret;
1778
1779 clean_x:
1780 Py_DECREF(ret->_internal_iter);
1781 array_iter_base_dealloc((PyArrayIterObject*)ret);
1782 PyArray_free((PyArrayObject*)ret);
1783 return NULL;
1784 }
1785
neighiter_dealloc(PyArrayNeighborhoodIterObject * iter)1786 static void neighiter_dealloc(PyArrayNeighborhoodIterObject* iter)
1787 {
1788 if (iter->mode == NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING) {
1789 if (PyArray_ISOBJECT(iter->_internal_iter->ao)) {
1790 Py_DECREF(*(PyObject**)iter->constant);
1791 }
1792 }
1793 PyDataMem_FREE(iter->constant);
1794 Py_DECREF(iter->_internal_iter);
1795
1796 array_iter_base_dealloc((PyArrayIterObject*)iter);
1797 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 };
1807