1 /*
2 * This file implements the CPython wrapper of NpyIter
3 *
4 * Copyright (c) 2010 by Mark Wiebe (mwwiebe@gmail.com)
5 * The University of British Columbia
6 *
7 * See LICENSE.txt for the license.
8 */
9 #define PY_SSIZE_T_CLEAN
10 #include "Python.h"
11 #include "structmember.h"
12
13 #define NPY_NO_DEPRECATED_API NPY_API_VERSION
14 #define _MULTIARRAYMODULE
15 #include <numpy/arrayobject.h>
16 #include "npy_config.h"
17 #include "npy_pycompat.h"
18 #include "alloc.h"
19 #include "common.h"
20 #include "conversion_utils.h"
21 #include "ctors.h"
22
23 /* Functions not part of the public NumPy C API */
24 npy_bool npyiter_has_writeback(NpyIter *iter);
25
26
27 typedef struct NewNpyArrayIterObject_tag NewNpyArrayIterObject;
28
29 struct NewNpyArrayIterObject_tag {
30 PyObject_HEAD
31 /* The iterator */
32 NpyIter *iter;
33 /* Flag indicating iteration started/stopped */
34 char started, finished;
35 /* Child to update for nested iteration */
36 NewNpyArrayIterObject *nested_child;
37 /* Cached values from the iterator */
38 NpyIter_IterNextFunc *iternext;
39 NpyIter_GetMultiIndexFunc *get_multi_index;
40 char **dataptrs;
41 PyArray_Descr **dtypes;
42 PyArrayObject **operands;
43 npy_intp *innerstrides, *innerloopsizeptr;
44 char readflags[NPY_MAXARGS];
45 char writeflags[NPY_MAXARGS];
46 };
47
npyiter_cache_values(NewNpyArrayIterObject * self)48 static int npyiter_cache_values(NewNpyArrayIterObject *self)
49 {
50 NpyIter *iter = self->iter;
51
52 /* iternext and get_multi_index functions */
53 self->iternext = NpyIter_GetIterNext(iter, NULL);
54 if (self->iternext == NULL) {
55 return -1;
56 }
57
58 if (NpyIter_HasMultiIndex(iter) && !NpyIter_HasDelayedBufAlloc(iter)) {
59 self->get_multi_index = NpyIter_GetGetMultiIndex(iter, NULL);
60 }
61 else {
62 self->get_multi_index = NULL;
63 }
64
65 /* Internal data pointers */
66 self->dataptrs = NpyIter_GetDataPtrArray(iter);
67 self->dtypes = NpyIter_GetDescrArray(iter);
68 self->operands = NpyIter_GetOperandArray(iter);
69
70 if (NpyIter_HasExternalLoop(iter)) {
71 self->innerstrides = NpyIter_GetInnerStrideArray(iter);
72 self->innerloopsizeptr = NpyIter_GetInnerLoopSizePtr(iter);
73 }
74 else {
75 self->innerstrides = NULL;
76 self->innerloopsizeptr = NULL;
77 }
78
79 /* The read/write settings */
80 NpyIter_GetReadFlags(iter, self->readflags);
81 NpyIter_GetWriteFlags(iter, self->writeflags);
82 return 0;
83 }
84
85 static PyObject *
npyiter_new(PyTypeObject * subtype,PyObject * NPY_UNUSED (args),PyObject * NPY_UNUSED (kwds))86 npyiter_new(PyTypeObject *subtype, PyObject *NPY_UNUSED(args),
87 PyObject *NPY_UNUSED(kwds))
88 {
89 NewNpyArrayIterObject *self;
90
91 self = (NewNpyArrayIterObject *)subtype->tp_alloc(subtype, 0);
92 if (self != NULL) {
93 self->iter = NULL;
94 self->nested_child = NULL;
95 }
96
97 return (PyObject *)self;
98 }
99
100 static int
NpyIter_GlobalFlagsConverter(PyObject * flags_in,npy_uint32 * flags)101 NpyIter_GlobalFlagsConverter(PyObject *flags_in, npy_uint32 *flags)
102 {
103 npy_uint32 tmpflags = 0;
104 int iflags, nflags;
105
106 PyObject *f;
107 char *str = NULL;
108 Py_ssize_t length = 0;
109 npy_uint32 flag;
110
111 if (flags_in == NULL || flags_in == Py_None) {
112 return 1;
113 }
114
115 if (!PyTuple_Check(flags_in) && !PyList_Check(flags_in)) {
116 PyErr_SetString(PyExc_ValueError,
117 "Iterator global flags must be a list or tuple of strings");
118 return 0;
119 }
120
121 nflags = PySequence_Size(flags_in);
122
123 for (iflags = 0; iflags < nflags; ++iflags) {
124 f = PySequence_GetItem(flags_in, iflags);
125 if (f == NULL) {
126 return 0;
127 }
128
129 if (PyUnicode_Check(f)) {
130 /* accept unicode input */
131 PyObject *f_str;
132 f_str = PyUnicode_AsASCIIString(f);
133 if (f_str == NULL) {
134 Py_DECREF(f);
135 return 0;
136 }
137 Py_DECREF(f);
138 f = f_str;
139 }
140
141 if (PyBytes_AsStringAndSize(f, &str, &length) < 0) {
142 Py_DECREF(f);
143 return 0;
144 }
145 /* Use switch statements to quickly isolate the right flag */
146 flag = 0;
147 switch (str[0]) {
148 case 'b':
149 if (strcmp(str, "buffered") == 0) {
150 flag = NPY_ITER_BUFFERED;
151 }
152 break;
153 case 'c':
154 if (length >= 6) switch (str[5]) {
155 case 'e':
156 if (strcmp(str, "c_index") == 0) {
157 flag = NPY_ITER_C_INDEX;
158 }
159 break;
160 case 'i':
161 if (strcmp(str, "copy_if_overlap") == 0) {
162 flag = NPY_ITER_COPY_IF_OVERLAP;
163 }
164 break;
165 case 'n':
166 if (strcmp(str, "common_dtype") == 0) {
167 flag = NPY_ITER_COMMON_DTYPE;
168 }
169 break;
170 }
171 break;
172 case 'd':
173 if (strcmp(str, "delay_bufalloc") == 0) {
174 flag = NPY_ITER_DELAY_BUFALLOC;
175 }
176 break;
177 case 'e':
178 if (strcmp(str, "external_loop") == 0) {
179 flag = NPY_ITER_EXTERNAL_LOOP;
180 }
181 break;
182 case 'f':
183 if (strcmp(str, "f_index") == 0) {
184 flag = NPY_ITER_F_INDEX;
185 }
186 break;
187 case 'g':
188 /*
189 * Documentation is grow_inner, but initial implementation
190 * was growinner, so allowing for either.
191 */
192 if (strcmp(str, "grow_inner") == 0 ||
193 strcmp(str, "growinner") == 0) {
194 flag = NPY_ITER_GROWINNER;
195 }
196 break;
197 case 'm':
198 if (strcmp(str, "multi_index") == 0) {
199 flag = NPY_ITER_MULTI_INDEX;
200 }
201 break;
202 case 'r':
203 if (strcmp(str, "ranged") == 0) {
204 flag = NPY_ITER_RANGED;
205 }
206 else if (strcmp(str, "refs_ok") == 0) {
207 flag = NPY_ITER_REFS_OK;
208 }
209 else if (strcmp(str, "reduce_ok") == 0) {
210 flag = NPY_ITER_REDUCE_OK;
211 }
212 break;
213 case 'z':
214 if (strcmp(str, "zerosize_ok") == 0) {
215 flag = NPY_ITER_ZEROSIZE_OK;
216 }
217 break;
218 }
219 if (flag == 0) {
220 PyErr_Format(PyExc_ValueError,
221 "Unexpected iterator global flag \"%s\"", str);
222 Py_DECREF(f);
223 return 0;
224 }
225 else {
226 tmpflags |= flag;
227 }
228 Py_DECREF(f);
229 }
230
231 *flags |= tmpflags;
232 return 1;
233 }
234
235 static int
NpyIter_OpFlagsConverter(PyObject * op_flags_in,npy_uint32 * op_flags)236 NpyIter_OpFlagsConverter(PyObject *op_flags_in,
237 npy_uint32 *op_flags)
238 {
239 int iflags, nflags;
240 npy_uint32 flag;
241
242 if (!PyTuple_Check(op_flags_in) && !PyList_Check(op_flags_in)) {
243 PyErr_SetString(PyExc_ValueError,
244 "op_flags must be a tuple or array of per-op flag-tuples");
245 return 0;
246 }
247
248 nflags = PySequence_Size(op_flags_in);
249
250 *op_flags = 0;
251 for (iflags = 0; iflags < nflags; ++iflags) {
252 PyObject *f;
253 char *str = NULL;
254 Py_ssize_t length = 0;
255
256 f = PySequence_GetItem(op_flags_in, iflags);
257 if (f == NULL) {
258 return 0;
259 }
260
261 if (PyUnicode_Check(f)) {
262 /* accept unicode input */
263 PyObject *f_str;
264 f_str = PyUnicode_AsASCIIString(f);
265 if (f_str == NULL) {
266 Py_DECREF(f);
267 return 0;
268 }
269 Py_DECREF(f);
270 f = f_str;
271 }
272
273 if (PyBytes_AsStringAndSize(f, &str, &length) < 0) {
274 PyErr_Clear();
275 Py_DECREF(f);
276 PyErr_SetString(PyExc_ValueError,
277 "op_flags must be a tuple or array of per-op flag-tuples");
278 return 0;
279 }
280
281 /* Use switch statements to quickly isolate the right flag */
282 flag = 0;
283 switch (str[0]) {
284 case 'a':
285 if (length > 2) switch(str[2]) {
286 case 'i':
287 if (strcmp(str, "aligned") == 0) {
288 flag = NPY_ITER_ALIGNED;
289 }
290 break;
291 case 'l':
292 if (strcmp(str, "allocate") == 0) {
293 flag = NPY_ITER_ALLOCATE;
294 }
295 break;
296 case 'r':
297 if (strcmp(str, "arraymask") == 0) {
298 flag = NPY_ITER_ARRAYMASK;
299 }
300 break;
301 }
302 break;
303 case 'c':
304 if (strcmp(str, "copy") == 0) {
305 flag = NPY_ITER_COPY;
306 }
307 if (strcmp(str, "contig") == 0) {
308 flag = NPY_ITER_CONTIG;
309 }
310 break;
311 case 'n':
312 switch (str[1]) {
313 case 'b':
314 if (strcmp(str, "nbo") == 0) {
315 flag = NPY_ITER_NBO;
316 }
317 break;
318 case 'o':
319 if (strcmp(str, "no_subtype") == 0) {
320 flag = NPY_ITER_NO_SUBTYPE;
321 }
322 else if (strcmp(str, "no_broadcast") == 0) {
323 flag = NPY_ITER_NO_BROADCAST;
324 }
325 break;
326 }
327 break;
328 case 'o':
329 if (strcmp(str, "overlap_assume_elementwise") == 0) {
330 flag = NPY_ITER_OVERLAP_ASSUME_ELEMENTWISE;
331 }
332 break;
333 case 'r':
334 if (length > 4) switch (str[4]) {
335 case 'o':
336 if (strcmp(str, "readonly") == 0) {
337 flag = NPY_ITER_READONLY;
338 }
339 break;
340 case 'w':
341 if (strcmp(str, "readwrite") == 0) {
342 flag = NPY_ITER_READWRITE;
343 }
344 break;
345 }
346 break;
347 case 'u':
348 switch (str[1]) {
349 case 'p':
350 if (strcmp(str, "updateifcopy") == 0) {
351 flag = NPY_ITER_UPDATEIFCOPY;
352 }
353 break;
354 }
355 break;
356 case 'v':
357 if (strcmp(str, "virtual") == 0) {
358 flag = NPY_ITER_VIRTUAL;
359 }
360 break;
361 case 'w':
362 if (length > 5) switch (str[5]) {
363 case 'o':
364 if (strcmp(str, "writeonly") == 0) {
365 flag = NPY_ITER_WRITEONLY;
366 }
367 break;
368 case 'm':
369 if (strcmp(str, "writemasked") == 0) {
370 flag = NPY_ITER_WRITEMASKED;
371 }
372 break;
373 }
374 break;
375 }
376 if (flag == 0) {
377 PyErr_Format(PyExc_ValueError,
378 "Unexpected per-op iterator flag \"%s\"", str);
379 Py_DECREF(f);
380 return 0;
381 }
382 else {
383 *op_flags |= flag;
384 }
385 Py_DECREF(f);
386 }
387
388 return 1;
389 }
390
391 static int
npyiter_convert_op_flags_array(PyObject * op_flags_in,npy_uint32 * op_flags_array,npy_intp nop)392 npyiter_convert_op_flags_array(PyObject *op_flags_in,
393 npy_uint32 *op_flags_array, npy_intp nop)
394 {
395 npy_intp iop;
396
397 if (!PyTuple_Check(op_flags_in) && !PyList_Check(op_flags_in)) {
398 PyErr_SetString(PyExc_ValueError,
399 "op_flags must be a tuple or array of per-op flag-tuples");
400 return 0;
401 }
402
403 if (PySequence_Size(op_flags_in) != nop) {
404 goto try_single_flags;
405 }
406
407 for (iop = 0; iop < nop; ++iop) {
408 PyObject *f = PySequence_GetItem(op_flags_in, iop);
409 if (f == NULL) {
410 return 0;
411 }
412 /* If the first item is a string, try as one set of flags */
413 if (iop == 0 && (PyBytes_Check(f) || PyUnicode_Check(f))) {
414 Py_DECREF(f);
415 goto try_single_flags;
416 }
417 if (NpyIter_OpFlagsConverter(f,
418 &op_flags_array[iop]) != 1) {
419 Py_DECREF(f);
420 return 0;
421 }
422
423 Py_DECREF(f);
424 }
425
426 return 1;
427
428 try_single_flags:
429 if (NpyIter_OpFlagsConverter(op_flags_in,
430 &op_flags_array[0]) != 1) {
431 return 0;
432 }
433
434 for (iop = 1; iop < nop; ++iop) {
435 op_flags_array[iop] = op_flags_array[0];
436 }
437
438 return 1;
439 }
440
441 static int
npyiter_convert_dtypes(PyObject * op_dtypes_in,PyArray_Descr ** op_dtypes,npy_intp nop)442 npyiter_convert_dtypes(PyObject *op_dtypes_in,
443 PyArray_Descr **op_dtypes,
444 npy_intp nop)
445 {
446 npy_intp iop;
447
448 /*
449 * If the input isn't a tuple of dtypes, try converting it as-is
450 * to a dtype, and replicating to all operands.
451 */
452 if ((!PyTuple_Check(op_dtypes_in) && !PyList_Check(op_dtypes_in)) ||
453 PySequence_Size(op_dtypes_in) != nop) {
454 goto try_single_dtype;
455 }
456
457 for (iop = 0; iop < nop; ++iop) {
458 PyObject *dtype = PySequence_GetItem(op_dtypes_in, iop);
459 if (dtype == NULL) {
460 npy_intp i;
461 for (i = 0; i < iop; ++i ) {
462 Py_XDECREF(op_dtypes[i]);
463 }
464 return 0;
465 }
466
467 /* Try converting the object to a descr */
468 if (PyArray_DescrConverter2(dtype, &op_dtypes[iop]) != 1) {
469 npy_intp i;
470 for (i = 0; i < iop; ++i ) {
471 Py_XDECREF(op_dtypes[i]);
472 }
473 Py_DECREF(dtype);
474 PyErr_Clear();
475 goto try_single_dtype;
476 }
477
478 Py_DECREF(dtype);
479 }
480
481 return 1;
482
483 try_single_dtype:
484 if (PyArray_DescrConverter2(op_dtypes_in, &op_dtypes[0]) == 1) {
485 for (iop = 1; iop < nop; ++iop) {
486 op_dtypes[iop] = op_dtypes[0];
487 Py_XINCREF(op_dtypes[iop]);
488 }
489 return 1;
490 }
491
492 return 0;
493 }
494
495 static int
npyiter_convert_op_axes(PyObject * op_axes_in,int nop,int ** op_axes,int * oa_ndim)496 npyiter_convert_op_axes(PyObject *op_axes_in, int nop,
497 int **op_axes, int *oa_ndim)
498 {
499 PyObject *a;
500 int iop;
501
502 if ((!PyTuple_Check(op_axes_in) && !PyList_Check(op_axes_in)) ||
503 PySequence_Size(op_axes_in) != nop) {
504 PyErr_SetString(PyExc_ValueError,
505 "op_axes must be a tuple/list matching the number of ops");
506 return 0;
507 }
508
509 *oa_ndim = -1;
510
511 /* Copy the tuples into op_axes */
512 for (iop = 0; iop < nop; ++iop) {
513 int idim;
514 a = PySequence_GetItem(op_axes_in, iop);
515 if (a == NULL) {
516 return 0;
517 }
518 if (a == Py_None) {
519 op_axes[iop] = NULL;
520 } else {
521 if (!PyTuple_Check(a) && !PyList_Check(a)) {
522 PyErr_SetString(PyExc_ValueError,
523 "Each entry of op_axes must be None "
524 "or a tuple/list");
525 Py_DECREF(a);
526 return 0;
527 }
528 if (*oa_ndim == -1) {
529 *oa_ndim = PySequence_Size(a);
530 if (*oa_ndim > NPY_MAXDIMS) {
531 PyErr_SetString(PyExc_ValueError,
532 "Too many dimensions in op_axes");
533 Py_DECREF(a);
534 return 0;
535 }
536 }
537 if (PySequence_Size(a) != *oa_ndim) {
538 PyErr_SetString(PyExc_ValueError,
539 "Each entry of op_axes must have the same size");
540 Py_DECREF(a);
541 return 0;
542 }
543 for (idim = 0; idim < *oa_ndim; ++idim) {
544 PyObject *v = PySequence_GetItem(a, idim);
545 if (v == NULL) {
546 Py_DECREF(a);
547 return 0;
548 }
549 /* numpy.newaxis is None */
550 if (v == Py_None) {
551 op_axes[iop][idim] = -1;
552 }
553 else {
554 op_axes[iop][idim] = PyArray_PyIntAsInt(v);
555 if (op_axes[iop][idim]==-1 &&
556 PyErr_Occurred()) {
557 Py_DECREF(a);
558 Py_DECREF(v);
559 return 0;
560 }
561 }
562 Py_DECREF(v);
563 }
564 }
565 Py_DECREF(a);
566 }
567
568 if (*oa_ndim == -1) {
569 PyErr_SetString(PyExc_ValueError,
570 "If op_axes is provided, at least one list of axes "
571 "must be contained within it");
572 return 0;
573 }
574
575 return 1;
576 }
577
578 /*
579 * Converts the operand array and op_flags array into the form
580 * NpyIter_AdvancedNew needs. Sets nop, and on success, each
581 * op[i] owns a reference to an array object.
582 */
583 static int
npyiter_convert_ops(PyObject * op_in,PyObject * op_flags_in,PyArrayObject ** op,npy_uint32 * op_flags,int * nop_out)584 npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
585 PyArrayObject **op, npy_uint32 *op_flags,
586 int *nop_out)
587 {
588 int iop, nop;
589
590 /* nop and op */
591 if (PyTuple_Check(op_in) || PyList_Check(op_in)) {
592 nop = PySequence_Size(op_in);
593 if (nop == 0) {
594 PyErr_SetString(PyExc_ValueError,
595 "Must provide at least one operand");
596 return 0;
597 }
598 if (nop > NPY_MAXARGS) {
599 PyErr_SetString(PyExc_ValueError, "Too many operands");
600 return 0;
601 }
602
603 for (iop = 0; iop < nop; ++iop) {
604 PyObject *item = PySequence_GetItem(op_in, iop);
605 if (item == NULL) {
606 npy_intp i;
607 for (i = 0; i < iop; ++i) {
608 Py_XDECREF(op[i]);
609 }
610 return 0;
611 }
612 else if (item == Py_None) {
613 Py_DECREF(item);
614 item = NULL;
615 }
616 /* This is converted to an array after op flags are retrieved */
617 op[iop] = (PyArrayObject *)item;
618 }
619 }
620 else {
621 nop = 1;
622 /* Is converted to an array after op flags are retrieved */
623 Py_INCREF(op_in);
624 op[0] = (PyArrayObject *)op_in;
625 }
626
627 *nop_out = nop;
628
629 /* op_flags */
630 if (op_flags_in == NULL || op_flags_in == Py_None) {
631 for (iop = 0; iop < nop; ++iop) {
632 /*
633 * By default, make NULL operands writeonly and flagged for
634 * allocation, and everything else readonly. To write
635 * to a provided operand, you must specify the write flag manually.
636 */
637 if (op[iop] == NULL) {
638 op_flags[iop] = NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE;
639 }
640 else {
641 op_flags[iop] = NPY_ITER_READONLY;
642 }
643 }
644 }
645 else if (npyiter_convert_op_flags_array(op_flags_in,
646 op_flags, nop) != 1) {
647 for (iop = 0; iop < nop; ++iop) {
648 Py_XDECREF(op[iop]);
649 }
650 *nop_out = 0;
651 return 0;
652 }
653
654 /* Now that we have the flags - convert all the ops to arrays */
655 for (iop = 0; iop < nop; ++iop) {
656 if (op[iop] != NULL) {
657 PyArrayObject *ao;
658 int fromanyflags = 0;
659
660 if (op_flags[iop]&(NPY_ITER_READWRITE|NPY_ITER_WRITEONLY)) {
661 fromanyflags |= NPY_ARRAY_WRITEBACKIFCOPY;
662 }
663 ao = (PyArrayObject *)PyArray_FROM_OF((PyObject *)op[iop],
664 fromanyflags);
665 if (ao == NULL) {
666 if (PyErr_Occurred() &&
667 PyErr_ExceptionMatches(PyExc_TypeError)) {
668 PyErr_SetString(PyExc_TypeError,
669 "Iterator operand is flagged as writeable, "
670 "but is an object which cannot be written "
671 "back to via WRITEBACKIFCOPY");
672 }
673 for (iop = 0; iop < nop; ++iop) {
674 Py_DECREF(op[iop]);
675 }
676 *nop_out = 0;
677 return 0;
678 }
679 Py_DECREF(op[iop]);
680 op[iop] = ao;
681 }
682 }
683
684 return 1;
685 }
686
687 static int
npyiter_init(NewNpyArrayIterObject * self,PyObject * args,PyObject * kwds)688 npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
689 {
690 static char *kwlist[] = {"op", "flags", "op_flags", "op_dtypes",
691 "order", "casting", "op_axes", "itershape",
692 "buffersize",
693 NULL};
694
695 PyObject *op_in = NULL, *op_flags_in = NULL,
696 *op_dtypes_in = NULL, *op_axes_in = NULL;
697
698 int iop, nop = 0;
699 PyArrayObject *op[NPY_MAXARGS];
700 npy_uint32 flags = 0;
701 NPY_ORDER order = NPY_KEEPORDER;
702 NPY_CASTING casting = NPY_SAFE_CASTING;
703 npy_uint32 op_flags[NPY_MAXARGS];
704 PyArray_Descr *op_request_dtypes[NPY_MAXARGS];
705 int oa_ndim = -1;
706 int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS];
707 int *op_axes[NPY_MAXARGS];
708 PyArray_Dims itershape = {NULL, -1};
709 int buffersize = 0;
710
711 if (self->iter != NULL) {
712 PyErr_SetString(PyExc_ValueError,
713 "Iterator was already initialized");
714 return -1;
715 }
716
717 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O&OOO&O&OO&i:nditer", kwlist,
718 &op_in,
719 NpyIter_GlobalFlagsConverter, &flags,
720 &op_flags_in,
721 &op_dtypes_in,
722 PyArray_OrderConverter, &order,
723 PyArray_CastingConverter, &casting,
724 &op_axes_in,
725 PyArray_OptionalIntpConverter, &itershape,
726 &buffersize)) {
727 npy_free_cache_dim_obj(itershape);
728 return -1;
729 }
730
731 /* Set the dtypes and ops to all NULL to start */
732 memset(op_request_dtypes, 0, sizeof(op_request_dtypes));
733
734 /* op and op_flags */
735 if (npyiter_convert_ops(op_in, op_flags_in, op, op_flags, &nop)
736 != 1) {
737 goto fail;
738 }
739
740 /* op_request_dtypes */
741 if (op_dtypes_in != NULL && op_dtypes_in != Py_None &&
742 npyiter_convert_dtypes(op_dtypes_in,
743 op_request_dtypes, nop) != 1) {
744 goto fail;
745 }
746
747 /* op_axes */
748 if (op_axes_in != NULL && op_axes_in != Py_None) {
749 /* Initialize to point to the op_axes arrays */
750 for (iop = 0; iop < nop; ++iop) {
751 op_axes[iop] = op_axes_arrays[iop];
752 }
753
754 if (npyiter_convert_op_axes(op_axes_in, nop,
755 op_axes, &oa_ndim) != 1) {
756 goto fail;
757 }
758 }
759
760 if (itershape.len != -1) {
761 if (oa_ndim == -1) {
762 oa_ndim = itershape.len;
763 memset(op_axes, 0, sizeof(op_axes[0]) * nop);
764 }
765 else if (oa_ndim != itershape.len) {
766 PyErr_SetString(PyExc_ValueError,
767 "'op_axes' and 'itershape' must have the same number "
768 "of entries equal to the iterator ndim");
769 goto fail;
770 }
771 }
772
773 self->iter = NpyIter_AdvancedNew(nop, op, flags, order, casting, op_flags,
774 op_request_dtypes,
775 oa_ndim, oa_ndim >= 0 ? op_axes : NULL,
776 itershape.ptr,
777 buffersize);
778
779 if (self->iter == NULL) {
780 goto fail;
781 }
782
783 /* Cache some values for the member functions to use */
784 if (npyiter_cache_values(self) < 0) {
785 goto fail;
786 }
787
788 if (NpyIter_GetIterSize(self->iter) == 0) {
789 self->started = 1;
790 self->finished = 1;
791 }
792 else {
793 self->started = 0;
794 self->finished = 0;
795 }
796
797 npy_free_cache_dim_obj(itershape);
798
799 /* Release the references we got to the ops and dtypes */
800 for (iop = 0; iop < nop; ++iop) {
801 Py_XDECREF(op[iop]);
802 Py_XDECREF(op_request_dtypes[iop]);
803 }
804
805 return 0;
806
807 fail:
808 npy_free_cache_dim_obj(itershape);
809 for (iop = 0; iop < nop; ++iop) {
810 Py_XDECREF(op[iop]);
811 Py_XDECREF(op_request_dtypes[iop]);
812 }
813 return -1;
814 }
815
816 NPY_NO_EXPORT PyObject *
NpyIter_NestedIters(PyObject * NPY_UNUSED (self),PyObject * args,PyObject * kwds)817 NpyIter_NestedIters(PyObject *NPY_UNUSED(self),
818 PyObject *args, PyObject *kwds)
819 {
820 static char *kwlist[] = {"op", "axes", "flags", "op_flags",
821 "op_dtypes", "order",
822 "casting", "buffersize",
823 NULL};
824
825 PyObject *op_in = NULL, *axes_in = NULL,
826 *op_flags_in = NULL, *op_dtypes_in = NULL;
827
828 int iop, nop = 0, inest, nnest = 0;
829 PyArrayObject *op[NPY_MAXARGS];
830 npy_uint32 flags = 0, flags_inner;
831 NPY_ORDER order = NPY_KEEPORDER;
832 NPY_CASTING casting = NPY_SAFE_CASTING;
833 npy_uint32 op_flags[NPY_MAXARGS], op_flags_inner[NPY_MAXARGS];
834 PyArray_Descr *op_request_dtypes[NPY_MAXARGS],
835 *op_request_dtypes_inner[NPY_MAXARGS];
836 int op_axes_data[NPY_MAXDIMS];
837 int *nested_op_axes[NPY_MAXDIMS];
838 int nested_naxes[NPY_MAXDIMS], iaxes, naxes;
839 int negones[NPY_MAXDIMS];
840 char used_axes[NPY_MAXDIMS];
841 int buffersize = 0;
842
843 PyObject *ret = NULL;
844
845 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|O&OOO&O&i", kwlist,
846 &op_in,
847 &axes_in,
848 NpyIter_GlobalFlagsConverter, &flags,
849 &op_flags_in,
850 &op_dtypes_in,
851 PyArray_OrderConverter, &order,
852 PyArray_CastingConverter, &casting,
853 &buffersize)) {
854 return NULL;
855 }
856
857 /* axes */
858 if (!PyTuple_Check(axes_in) && !PyList_Check(axes_in)) {
859 PyErr_SetString(PyExc_ValueError,
860 "axes must be a tuple of axis arrays");
861 return NULL;
862 }
863 nnest = PySequence_Size(axes_in);
864 if (nnest < 2) {
865 PyErr_SetString(PyExc_ValueError,
866 "axes must have at least 2 entries for nested iteration");
867 return NULL;
868 }
869 naxes = 0;
870 memset(used_axes, 0, NPY_MAXDIMS);
871 for (inest = 0; inest < nnest; ++inest) {
872 PyObject *item = PySequence_GetItem(axes_in, inest);
873 npy_intp i;
874 if (item == NULL) {
875 return NULL;
876 }
877 if (!PyTuple_Check(item) && !PyList_Check(item)) {
878 PyErr_SetString(PyExc_ValueError,
879 "Each item in axes must be a an integer tuple");
880 Py_DECREF(item);
881 return NULL;
882 }
883 nested_naxes[inest] = PySequence_Size(item);
884 if (naxes + nested_naxes[inest] > NPY_MAXDIMS) {
885 PyErr_SetString(PyExc_ValueError,
886 "Too many axes given");
887 Py_DECREF(item);
888 return NULL;
889 }
890 for (i = 0; i < nested_naxes[inest]; ++i) {
891 PyObject *v = PySequence_GetItem(item, i);
892 npy_intp axis;
893 if (v == NULL) {
894 Py_DECREF(item);
895 return NULL;
896 }
897 axis = PyLong_AsLong(v);
898 Py_DECREF(v);
899 if (axis < 0 || axis >= NPY_MAXDIMS) {
900 PyErr_SetString(PyExc_ValueError,
901 "An axis is out of bounds");
902 Py_DECREF(item);
903 return NULL;
904 }
905 /*
906 * This check is very important, without it out of bounds
907 * data accesses are possible.
908 */
909 if (used_axes[axis] != 0) {
910 PyErr_SetString(PyExc_ValueError,
911 "An axis is used more than once");
912 Py_DECREF(item);
913 return NULL;
914 }
915 used_axes[axis] = 1;
916 op_axes_data[naxes+i] = axis;
917 }
918 nested_op_axes[inest] = &op_axes_data[naxes];
919 naxes += nested_naxes[inest];
920 Py_DECREF(item);
921 }
922
923 /* op and op_flags */
924 if (npyiter_convert_ops(op_in, op_flags_in, op, op_flags, &nop)
925 != 1) {
926 return NULL;
927 }
928
929 /* Set the dtypes to all NULL to start as well */
930 memset(op_request_dtypes, 0, sizeof(op_request_dtypes[0])*nop);
931 memset(op_request_dtypes_inner, 0,
932 sizeof(op_request_dtypes_inner[0])*nop);
933
934 /* op_request_dtypes */
935 if (op_dtypes_in != NULL && op_dtypes_in != Py_None &&
936 npyiter_convert_dtypes(op_dtypes_in,
937 op_request_dtypes, nop) != 1) {
938 goto fail;
939 }
940
941 ret = PyTuple_New(nnest);
942 if (ret == NULL) {
943 goto fail;
944 }
945
946 /* For broadcasting allocated arrays */
947 for (iaxes = 0; iaxes < naxes; ++iaxes) {
948 negones[iaxes] = -1;
949 }
950
951 /*
952 * Clear any unnecessary ALLOCATE flags, so we can use them
953 * to indicate exactly the allocated outputs. Also, separate
954 * the inner loop flags.
955 */
956 for (iop = 0; iop < nop; ++iop) {
957 if ((op_flags[iop]&NPY_ITER_ALLOCATE) && op[iop] != NULL) {
958 op_flags[iop] &= ~NPY_ITER_ALLOCATE;
959 }
960
961 /*
962 * Clear any flags allowing copies or output allocation for
963 * the inner loop.
964 */
965 op_flags_inner[iop] = op_flags[iop] & ~(NPY_ITER_COPY|
966 NPY_ITER_UPDATEIFCOPY|
967 NPY_ITER_ALLOCATE);
968 /*
969 * If buffering is enabled and copying is not,
970 * clear the nbo_aligned flag and strip the data type
971 * for the outer loops.
972 */
973 if ((flags&(NPY_ITER_BUFFERED)) &&
974 !(op_flags[iop]&(NPY_ITER_COPY|
975 NPY_ITER_UPDATEIFCOPY|
976 NPY_ITER_ALLOCATE))) {
977 op_flags[iop] &= ~(NPY_ITER_NBO|NPY_ITER_ALIGNED|NPY_ITER_CONTIG);
978 op_request_dtypes_inner[iop] = op_request_dtypes[iop];
979 op_request_dtypes[iop] = NULL;
980 }
981 }
982
983 /* Only the inner loop gets the buffering and no inner flags */
984 flags_inner = flags&~NPY_ITER_COMMON_DTYPE;
985 flags &= ~(NPY_ITER_EXTERNAL_LOOP|
986 NPY_ITER_BUFFERED);
987
988 for (inest = 0; inest < nnest; ++inest) {
989 NewNpyArrayIterObject *iter;
990 int *op_axes_nop[NPY_MAXARGS];
991
992 /*
993 * All the operands' op_axes are the same, except for
994 * allocated outputs.
995 */
996 for (iop = 0; iop < nop; ++iop) {
997 if (op_flags[iop]&NPY_ITER_ALLOCATE) {
998 if (inest == 0) {
999 op_axes_nop[iop] = NULL;
1000 }
1001 else {
1002 op_axes_nop[iop] = negones;
1003 }
1004 }
1005 else {
1006 op_axes_nop[iop] = nested_op_axes[inest];
1007 }
1008 }
1009
1010 /*
1011 printf("\n");
1012 for (iop = 0; iop < nop; ++iop) {
1013 npy_intp i;
1014
1015 for (i = 0; i < nested_naxes[inest]; ++i) {
1016 printf("%d ", (int)op_axes_nop[iop][i]);
1017 }
1018 printf("\n");
1019 }
1020 */
1021
1022 /* Allocate the iterator */
1023 iter = (NewNpyArrayIterObject *)npyiter_new(&NpyIter_Type, NULL, NULL);
1024 if (iter == NULL) {
1025 Py_DECREF(ret);
1026 goto fail;
1027 }
1028
1029 if (inest < nnest-1) {
1030 iter->iter = NpyIter_AdvancedNew(nop, op, flags, order,
1031 casting, op_flags, op_request_dtypes,
1032 nested_naxes[inest], op_axes_nop,
1033 NULL,
1034 0);
1035 }
1036 else {
1037 iter->iter = NpyIter_AdvancedNew(nop, op, flags_inner, order,
1038 casting, op_flags_inner,
1039 op_request_dtypes_inner,
1040 nested_naxes[inest], op_axes_nop,
1041 NULL,
1042 buffersize);
1043 }
1044
1045 if (iter->iter == NULL) {
1046 Py_DECREF(ret);
1047 goto fail;
1048 }
1049
1050 /* Cache some values for the member functions to use */
1051 if (npyiter_cache_values(iter) < 0) {
1052 Py_DECREF(ret);
1053 goto fail;
1054 }
1055
1056 if (NpyIter_GetIterSize(iter->iter) == 0) {
1057 iter->started = 1;
1058 iter->finished = 1;
1059 }
1060 else {
1061 iter->started = 0;
1062 iter->finished = 0;
1063 }
1064
1065 /*
1066 * If there are any allocated outputs or any copies were made,
1067 * adjust op so that the other iterators use the same ones.
1068 */
1069 if (inest == 0) {
1070 PyArrayObject **operands = NpyIter_GetOperandArray(iter->iter);
1071 for (iop = 0; iop < nop; ++iop) {
1072 if (op[iop] != operands[iop]) {
1073 Py_XDECREF(op[iop]);
1074 op[iop] = operands[iop];
1075 Py_INCREF(op[iop]);
1076 }
1077
1078 /*
1079 * Clear any flags allowing copies for
1080 * the rest of the iterators
1081 */
1082 op_flags[iop] &= ~(NPY_ITER_COPY|
1083 NPY_ITER_UPDATEIFCOPY);
1084 }
1085 /* Clear the common dtype flag for the rest of the iterators */
1086 flags &= ~NPY_ITER_COMMON_DTYPE;
1087 }
1088
1089 PyTuple_SET_ITEM(ret, inest, (PyObject *)iter);
1090 }
1091
1092 /* Release our references to the ops and dtypes */
1093 for (iop = 0; iop < nop; ++iop) {
1094 Py_XDECREF(op[iop]);
1095 Py_XDECREF(op_request_dtypes[iop]);
1096 Py_XDECREF(op_request_dtypes_inner[iop]);
1097 }
1098
1099 /* Set up the nested child references */
1100 for (inest = 0; inest < nnest-1; ++inest) {
1101 NewNpyArrayIterObject *iter;
1102 iter = (NewNpyArrayIterObject *)PyTuple_GET_ITEM(ret, inest);
1103 /*
1104 * Indicates which iterator to reset with new base pointers
1105 * each iteration step.
1106 */
1107 iter->nested_child =
1108 (NewNpyArrayIterObject *)PyTuple_GET_ITEM(ret, inest+1);
1109 Py_INCREF(iter->nested_child);
1110 /*
1111 * Need to do a nested reset so all the iterators point
1112 * at the right data
1113 */
1114 if (NpyIter_ResetBasePointers(iter->nested_child->iter,
1115 iter->dataptrs, NULL) != NPY_SUCCEED) {
1116 Py_DECREF(ret);
1117 return NULL;
1118 }
1119 }
1120
1121 return ret;
1122
1123 fail:
1124 for (iop = 0; iop < nop; ++iop) {
1125 Py_XDECREF(op[iop]);
1126 Py_XDECREF(op_request_dtypes[iop]);
1127 Py_XDECREF(op_request_dtypes_inner[iop]);
1128 }
1129 return NULL;
1130 }
1131
1132
1133 static void
npyiter_dealloc(NewNpyArrayIterObject * self)1134 npyiter_dealloc(NewNpyArrayIterObject *self)
1135 {
1136 if (self->iter) {
1137 if (npyiter_has_writeback(self->iter)) {
1138 if (PyErr_WarnEx(PyExc_RuntimeWarning,
1139 "Temporary data has not been written back to one of the "
1140 "operands. Typically nditer is used as a context manager "
1141 "otherwise 'close' must be called before reading iteration "
1142 "results.", 1) < 0) {
1143 PyObject *s;
1144
1145 s = PyUnicode_FromString("npyiter_dealloc");
1146 if (s) {
1147 PyErr_WriteUnraisable(s);
1148 Py_DECREF(s);
1149 }
1150 else {
1151 PyErr_WriteUnraisable(Py_None);
1152 }
1153 }
1154 }
1155 NpyIter_Deallocate(self->iter);
1156 self->iter = NULL;
1157 Py_XDECREF(self->nested_child);
1158 self->nested_child = NULL;
1159 }
1160 Py_TYPE(self)->tp_free((PyObject*)self);
1161 }
1162
1163 static int
npyiter_resetbasepointers(NewNpyArrayIterObject * self)1164 npyiter_resetbasepointers(NewNpyArrayIterObject *self)
1165 {
1166 while (self->nested_child) {
1167 if (NpyIter_ResetBasePointers(self->nested_child->iter,
1168 self->dataptrs, NULL) != NPY_SUCCEED) {
1169 return NPY_FAIL;
1170 }
1171 self = self->nested_child;
1172 if (NpyIter_GetIterSize(self->iter) == 0) {
1173 self->started = 1;
1174 self->finished = 1;
1175 }
1176 else {
1177 self->started = 0;
1178 self->finished = 0;
1179 }
1180 }
1181
1182 return NPY_SUCCEED;
1183 }
1184
1185 static PyObject *
npyiter_reset(NewNpyArrayIterObject * self)1186 npyiter_reset(NewNpyArrayIterObject *self)
1187 {
1188 if (self->iter == NULL) {
1189 PyErr_SetString(PyExc_ValueError,
1190 "Iterator is invalid");
1191 return NULL;
1192 }
1193
1194 if (NpyIter_Reset(self->iter, NULL) != NPY_SUCCEED) {
1195 return NULL;
1196 }
1197 if (NpyIter_GetIterSize(self->iter) == 0) {
1198 self->started = 1;
1199 self->finished = 1;
1200 }
1201 else {
1202 self->started = 0;
1203 self->finished = 0;
1204 }
1205
1206 if (self->get_multi_index == NULL && NpyIter_HasMultiIndex(self->iter)) {
1207 self->get_multi_index = NpyIter_GetGetMultiIndex(self->iter, NULL);
1208 }
1209
1210 /* If there is nesting, the nested iterators should be reset */
1211 if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1212 return NULL;
1213 }
1214
1215 Py_RETURN_NONE;
1216 }
1217
1218 /*
1219 * Makes a copy of the iterator. Note that the nesting is not
1220 * copied.
1221 */
1222 static PyObject *
npyiter_copy(NewNpyArrayIterObject * self)1223 npyiter_copy(NewNpyArrayIterObject *self)
1224 {
1225 NewNpyArrayIterObject *iter;
1226
1227 if (self->iter == NULL) {
1228 PyErr_SetString(PyExc_ValueError,
1229 "Iterator is invalid");
1230 return NULL;
1231 }
1232
1233 /* Allocate the iterator */
1234 iter = (NewNpyArrayIterObject *)npyiter_new(&NpyIter_Type, NULL, NULL);
1235 if (iter == NULL) {
1236 return NULL;
1237 }
1238
1239 /* Copy the C iterator */
1240 iter->iter = NpyIter_Copy(self->iter);
1241 if (iter->iter == NULL) {
1242 Py_DECREF(iter);
1243 return NULL;
1244 }
1245
1246 /* Cache some values for the member functions to use */
1247 if (npyiter_cache_values(iter) < 0) {
1248 Py_DECREF(iter);
1249 return NULL;
1250 }
1251
1252 iter->started = self->started;
1253 iter->finished = self->finished;
1254
1255 return (PyObject *)iter;
1256 }
1257
1258 static PyObject *
npyiter_iternext(NewNpyArrayIterObject * self)1259 npyiter_iternext(NewNpyArrayIterObject *self)
1260 {
1261 if (self->iter != NULL && self->iternext != NULL &&
1262 !self->finished && self->iternext(self->iter)) {
1263 /* If there is nesting, the nested iterators should be reset */
1264 if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1265 return NULL;
1266 }
1267
1268 Py_RETURN_TRUE;
1269 }
1270 else {
1271 if (PyErr_Occurred()) {
1272 /* casting error, buffer cleanup will occur at reset or dealloc */
1273 return NULL;
1274 }
1275 self->finished = 1;
1276 Py_RETURN_FALSE;
1277 }
1278 }
1279
1280 static PyObject *
npyiter_remove_axis(NewNpyArrayIterObject * self,PyObject * args)1281 npyiter_remove_axis(NewNpyArrayIterObject *self, PyObject *args)
1282 {
1283 int axis = 0;
1284
1285 if (self->iter == NULL) {
1286 PyErr_SetString(PyExc_ValueError,
1287 "Iterator is invalid");
1288 return NULL;
1289 }
1290
1291 if (!PyArg_ParseTuple(args, "i:remove_axis", &axis)) {
1292 return NULL;
1293 }
1294
1295 if (NpyIter_RemoveAxis(self->iter, axis) != NPY_SUCCEED) {
1296 return NULL;
1297 }
1298 /* RemoveAxis invalidates cached values */
1299 if (npyiter_cache_values(self) < 0) {
1300 return NULL;
1301 }
1302 /* RemoveAxis also resets the iterator */
1303 if (NpyIter_GetIterSize(self->iter) == 0) {
1304 self->started = 1;
1305 self->finished = 1;
1306 }
1307 else {
1308 self->started = 0;
1309 self->finished = 0;
1310 }
1311
1312 Py_RETURN_NONE;
1313 }
1314
1315 static PyObject *
npyiter_remove_multi_index(NewNpyArrayIterObject * self)1316 npyiter_remove_multi_index(NewNpyArrayIterObject *self)
1317 {
1318 if (self->iter == NULL) {
1319 PyErr_SetString(PyExc_ValueError,
1320 "Iterator is invalid");
1321 return NULL;
1322 }
1323
1324 NpyIter_RemoveMultiIndex(self->iter);
1325 /* RemoveMultiIndex invalidates cached values */
1326 npyiter_cache_values(self);
1327 /* RemoveMultiIndex also resets the iterator */
1328 if (NpyIter_GetIterSize(self->iter) == 0) {
1329 self->started = 1;
1330 self->finished = 1;
1331 }
1332 else {
1333 self->started = 0;
1334 self->finished = 0;
1335 }
1336
1337 Py_RETURN_NONE;
1338 }
1339
1340 static PyObject *
npyiter_enable_external_loop(NewNpyArrayIterObject * self)1341 npyiter_enable_external_loop(NewNpyArrayIterObject *self)
1342 {
1343 if (self->iter == NULL) {
1344 PyErr_SetString(PyExc_ValueError,
1345 "Iterator is invalid");
1346 return NULL;
1347 }
1348
1349 NpyIter_EnableExternalLoop(self->iter);
1350 /* EnableExternalLoop invalidates cached values */
1351 npyiter_cache_values(self);
1352 /* EnableExternalLoop also resets the iterator */
1353 if (NpyIter_GetIterSize(self->iter) == 0) {
1354 self->started = 1;
1355 self->finished = 1;
1356 }
1357 else {
1358 self->started = 0;
1359 self->finished = 0;
1360 }
1361
1362 Py_RETURN_NONE;
1363 }
1364
1365 static PyObject *
npyiter_debug_print(NewNpyArrayIterObject * self)1366 npyiter_debug_print(NewNpyArrayIterObject *self)
1367 {
1368 if (self->iter != NULL) {
1369 NpyIter_DebugPrint(self->iter);
1370 }
1371 else {
1372 printf("Iterator: (nil)\n");
1373 }
1374
1375 Py_RETURN_NONE;
1376 }
1377
1378 NPY_NO_EXPORT PyObject *
1379 npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i);
1380
npyiter_value_get(NewNpyArrayIterObject * self)1381 static PyObject *npyiter_value_get(NewNpyArrayIterObject *self)
1382 {
1383 PyObject *ret;
1384
1385 npy_intp iop, nop;
1386
1387 if (self->iter == NULL || self->finished) {
1388 PyErr_SetString(PyExc_ValueError,
1389 "Iterator is past the end");
1390 return NULL;
1391 }
1392
1393 nop = NpyIter_GetNOp(self->iter);
1394
1395 /* Return an array or tuple of arrays with the values */
1396 if (nop == 1) {
1397 ret = npyiter_seq_item(self, 0);
1398 }
1399 else {
1400 ret = PyTuple_New(nop);
1401 if (ret == NULL) {
1402 return NULL;
1403 }
1404 for (iop = 0; iop < nop; ++iop) {
1405 PyObject *a = npyiter_seq_item(self, iop);
1406 if (a == NULL) {
1407 Py_DECREF(ret);
1408 return NULL;
1409 }
1410 PyTuple_SET_ITEM(ret, iop, a);
1411 }
1412 }
1413
1414 return ret;
1415 }
1416
npyiter_operands_get(NewNpyArrayIterObject * self)1417 static PyObject *npyiter_operands_get(NewNpyArrayIterObject *self)
1418 {
1419 PyObject *ret;
1420
1421 npy_intp iop, nop;
1422 PyArrayObject **operands;
1423
1424 if (self->iter == NULL) {
1425 PyErr_SetString(PyExc_ValueError,
1426 "Iterator is invalid");
1427 return NULL;
1428 }
1429 nop = NpyIter_GetNOp(self->iter);
1430 operands = self->operands;
1431
1432 ret = PyTuple_New(nop);
1433 if (ret == NULL) {
1434 return NULL;
1435 }
1436 for (iop = 0; iop < nop; ++iop) {
1437 PyObject *operand = (PyObject *)operands[iop];
1438
1439 Py_INCREF(operand);
1440 PyTuple_SET_ITEM(ret, iop, operand);
1441 }
1442
1443 return ret;
1444 }
1445
npyiter_itviews_get(NewNpyArrayIterObject * self)1446 static PyObject *npyiter_itviews_get(NewNpyArrayIterObject *self)
1447 {
1448 PyObject *ret;
1449
1450 npy_intp iop, nop;
1451
1452 if (self->iter == NULL) {
1453 PyErr_SetString(PyExc_ValueError,
1454 "Iterator is invalid");
1455 return NULL;
1456 }
1457 nop = NpyIter_GetNOp(self->iter);
1458
1459 ret = PyTuple_New(nop);
1460 if (ret == NULL) {
1461 return NULL;
1462 }
1463 for (iop = 0; iop < nop; ++iop) {
1464 PyArrayObject *view = NpyIter_GetIterView(self->iter, iop);
1465
1466 if (view == NULL) {
1467 Py_DECREF(ret);
1468 return NULL;
1469 }
1470 PyTuple_SET_ITEM(ret, iop, (PyObject *)view);
1471 }
1472
1473 return ret;
1474 }
1475
1476 static PyObject *
npyiter_next(NewNpyArrayIterObject * self)1477 npyiter_next(NewNpyArrayIterObject *self)
1478 {
1479 if (self->iter == NULL || self->iternext == NULL ||
1480 self->finished) {
1481 return NULL;
1482 }
1483
1484 /*
1485 * Use the started flag for the Python iteration protocol to work
1486 * when buffering is enabled.
1487 */
1488 if (self->started) {
1489 if (!self->iternext(self->iter)) {
1490 /*
1491 * A casting error may be set here (or no error causing a
1492 * StopIteration). Buffers may only be cleaned up later.
1493 */
1494 self->finished = 1;
1495 return NULL;
1496 }
1497
1498 /* If there is nesting, the nested iterators should be reset */
1499 if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1500 return NULL;
1501 }
1502 }
1503 self->started = 1;
1504
1505 return npyiter_value_get(self);
1506 };
1507
npyiter_shape_get(NewNpyArrayIterObject * self)1508 static PyObject *npyiter_shape_get(NewNpyArrayIterObject *self)
1509 {
1510 PyObject *ret;
1511 npy_intp idim, ndim, shape[NPY_MAXDIMS];
1512
1513 if (self->iter == NULL || self->finished) {
1514 PyErr_SetString(PyExc_ValueError,
1515 "Iterator is past the end");
1516 return NULL;
1517 }
1518
1519 if (NpyIter_GetShape(self->iter, shape) == NPY_SUCCEED) {
1520 ndim = NpyIter_GetNDim(self->iter);
1521 ret = PyTuple_New(ndim);
1522 if (ret != NULL) {
1523 for (idim = 0; idim < ndim; ++idim) {
1524 PyTuple_SET_ITEM(ret, idim,
1525 PyLong_FromLong(shape[idim]));
1526 }
1527 return ret;
1528 }
1529 }
1530
1531 return NULL;
1532 }
1533
npyiter_multi_index_get(NewNpyArrayIterObject * self)1534 static PyObject *npyiter_multi_index_get(NewNpyArrayIterObject *self)
1535 {
1536 PyObject *ret;
1537 npy_intp idim, ndim, multi_index[NPY_MAXDIMS];
1538
1539 if (self->iter == NULL || self->finished) {
1540 PyErr_SetString(PyExc_ValueError,
1541 "Iterator is past the end");
1542 return NULL;
1543 }
1544
1545 if (self->get_multi_index != NULL) {
1546 ndim = NpyIter_GetNDim(self->iter);
1547 self->get_multi_index(self->iter, multi_index);
1548 ret = PyTuple_New(ndim);
1549 if (ret == NULL) {
1550 return NULL;
1551 }
1552 for (idim = 0; idim < ndim; ++idim) {
1553 PyTuple_SET_ITEM(ret, idim,
1554 PyLong_FromLong(multi_index[idim]));
1555 }
1556 return ret;
1557 }
1558 else {
1559 if (!NpyIter_HasMultiIndex(self->iter)) {
1560 PyErr_SetString(PyExc_ValueError,
1561 "Iterator is not tracking a multi-index");
1562 return NULL;
1563 }
1564 else if (NpyIter_HasDelayedBufAlloc(self->iter)) {
1565 PyErr_SetString(PyExc_ValueError,
1566 "Iterator construction used delayed buffer allocation, "
1567 "and no reset has been done yet");
1568 return NULL;
1569 }
1570 else {
1571 PyErr_SetString(PyExc_ValueError,
1572 "Iterator is in an invalid state");
1573 return NULL;
1574 }
1575 }
1576 }
1577
1578 static int
npyiter_multi_index_set(NewNpyArrayIterObject * self,PyObject * value)1579 npyiter_multi_index_set(NewNpyArrayIterObject *self, PyObject *value)
1580 {
1581 npy_intp idim, ndim, multi_index[NPY_MAXDIMS];
1582
1583 if (value == NULL) {
1584 PyErr_SetString(PyExc_AttributeError,
1585 "Cannot delete nditer multi_index");
1586 return -1;
1587 }
1588 if (self->iter == NULL) {
1589 PyErr_SetString(PyExc_ValueError,
1590 "Iterator is invalid");
1591 return -1;
1592 }
1593
1594 if (NpyIter_HasMultiIndex(self->iter)) {
1595 ndim = NpyIter_GetNDim(self->iter);
1596 if (!PySequence_Check(value)) {
1597 PyErr_SetString(PyExc_ValueError,
1598 "multi_index must be set with a sequence");
1599 return -1;
1600 }
1601 if (PySequence_Size(value) != ndim) {
1602 PyErr_SetString(PyExc_ValueError,
1603 "Wrong number of indices");
1604 return -1;
1605 }
1606 for (idim = 0; idim < ndim; ++idim) {
1607 PyObject *v = PySequence_GetItem(value, idim);
1608 multi_index[idim] = PyLong_AsLong(v);
1609 if (error_converting(multi_index[idim])) {
1610 Py_XDECREF(v);
1611 return -1;
1612 }
1613 }
1614 if (NpyIter_GotoMultiIndex(self->iter, multi_index) != NPY_SUCCEED) {
1615 return -1;
1616 }
1617 self->started = 0;
1618 self->finished = 0;
1619
1620 /* If there is nesting, the nested iterators should be reset */
1621 if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1622 return -1;
1623 }
1624
1625 return 0;
1626 }
1627 else {
1628 PyErr_SetString(PyExc_ValueError,
1629 "Iterator is not tracking a multi-index");
1630 return -1;
1631 }
1632 }
1633
npyiter_index_get(NewNpyArrayIterObject * self)1634 static PyObject *npyiter_index_get(NewNpyArrayIterObject *self)
1635 {
1636 if (self->iter == NULL || self->finished) {
1637 PyErr_SetString(PyExc_ValueError,
1638 "Iterator is past the end");
1639 return NULL;
1640 }
1641
1642 if (NpyIter_HasIndex(self->iter)) {
1643 npy_intp ind = *NpyIter_GetIndexPtr(self->iter);
1644 return PyLong_FromLong(ind);
1645 }
1646 else {
1647 PyErr_SetString(PyExc_ValueError,
1648 "Iterator does not have an index");
1649 return NULL;
1650 }
1651 }
1652
npyiter_index_set(NewNpyArrayIterObject * self,PyObject * value)1653 static int npyiter_index_set(NewNpyArrayIterObject *self, PyObject *value)
1654 {
1655 if (value == NULL) {
1656 PyErr_SetString(PyExc_AttributeError,
1657 "Cannot delete nditer index");
1658 return -1;
1659 }
1660 if (self->iter == NULL) {
1661 PyErr_SetString(PyExc_ValueError,
1662 "Iterator is invalid");
1663 return -1;
1664 }
1665
1666 if (NpyIter_HasIndex(self->iter)) {
1667 npy_intp ind;
1668 ind = PyLong_AsLong(value);
1669 if (error_converting(ind)) {
1670 return -1;
1671 }
1672 if (NpyIter_GotoIndex(self->iter, ind) != NPY_SUCCEED) {
1673 return -1;
1674 }
1675 self->started = 0;
1676 self->finished = 0;
1677
1678 /* If there is nesting, the nested iterators should be reset */
1679 if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1680 return -1;
1681 }
1682
1683 return 0;
1684 }
1685 else {
1686 PyErr_SetString(PyExc_ValueError,
1687 "Iterator does not have an index");
1688 return -1;
1689 }
1690 }
1691
npyiter_iterindex_get(NewNpyArrayIterObject * self)1692 static PyObject *npyiter_iterindex_get(NewNpyArrayIterObject *self)
1693 {
1694 if (self->iter == NULL || self->finished) {
1695 PyErr_SetString(PyExc_ValueError,
1696 "Iterator is past the end");
1697 return NULL;
1698 }
1699
1700 return PyLong_FromLong(NpyIter_GetIterIndex(self->iter));
1701 }
1702
npyiter_iterindex_set(NewNpyArrayIterObject * self,PyObject * value)1703 static int npyiter_iterindex_set(NewNpyArrayIterObject *self, PyObject *value)
1704 {
1705 npy_intp iterindex;
1706
1707 if (value == NULL) {
1708 PyErr_SetString(PyExc_AttributeError,
1709 "Cannot delete nditer iterindex");
1710 return -1;
1711 }
1712 if (self->iter == NULL) {
1713 PyErr_SetString(PyExc_ValueError,
1714 "Iterator is invalid");
1715 return -1;
1716 }
1717
1718 iterindex = PyLong_AsLong(value);
1719 if (error_converting(iterindex)) {
1720 return -1;
1721 }
1722 if (NpyIter_GotoIterIndex(self->iter, iterindex) != NPY_SUCCEED) {
1723 return -1;
1724 }
1725 self->started = 0;
1726 self->finished = 0;
1727
1728 /* If there is nesting, the nested iterators should be reset */
1729 if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1730 return -1;
1731 }
1732
1733 return 0;
1734 }
1735
npyiter_iterrange_get(NewNpyArrayIterObject * self)1736 static PyObject *npyiter_iterrange_get(NewNpyArrayIterObject *self)
1737 {
1738 npy_intp istart = 0, iend = 0;
1739 PyObject *ret;
1740
1741 if (self->iter == NULL) {
1742 PyErr_SetString(PyExc_ValueError,
1743 "Iterator is invalid");
1744 return NULL;
1745 }
1746
1747 NpyIter_GetIterIndexRange(self->iter, &istart, &iend);
1748
1749 ret = PyTuple_New(2);
1750 if (ret == NULL) {
1751 return NULL;
1752 }
1753
1754 PyTuple_SET_ITEM(ret, 0, PyLong_FromLong(istart));
1755 PyTuple_SET_ITEM(ret, 1, PyLong_FromLong(iend));
1756
1757 return ret;
1758 }
1759
npyiter_iterrange_set(NewNpyArrayIterObject * self,PyObject * value)1760 static int npyiter_iterrange_set(NewNpyArrayIterObject *self, PyObject *value)
1761 {
1762 npy_intp istart = 0, iend = 0;
1763
1764 if (value == NULL) {
1765 PyErr_SetString(PyExc_AttributeError,
1766 "Cannot delete nditer iterrange");
1767 return -1;
1768 }
1769 if (self->iter == NULL) {
1770 PyErr_SetString(PyExc_ValueError,
1771 "Iterator is invalid");
1772 return -1;
1773 }
1774
1775 if (!PyArg_ParseTuple(value, "nn", &istart, &iend)) {
1776 return -1;
1777 }
1778
1779 if (NpyIter_ResetToIterIndexRange(self->iter, istart, iend, NULL)
1780 != NPY_SUCCEED) {
1781 return -1;
1782 }
1783 if (istart < iend) {
1784 self->started = self->finished = 0;
1785 }
1786 else {
1787 self->started = self->finished = 1;
1788 }
1789
1790 if (self->get_multi_index == NULL && NpyIter_HasMultiIndex(self->iter)) {
1791 self->get_multi_index = NpyIter_GetGetMultiIndex(self->iter, NULL);
1792 }
1793
1794 /* If there is nesting, the nested iterators should be reset */
1795 if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1796 return -1;
1797 }
1798
1799 return 0;
1800 }
1801
npyiter_has_delayed_bufalloc_get(NewNpyArrayIterObject * self)1802 static PyObject *npyiter_has_delayed_bufalloc_get(NewNpyArrayIterObject *self)
1803 {
1804 if (self->iter == NULL) {
1805 PyErr_SetString(PyExc_ValueError,
1806 "Iterator is invalid");
1807 return NULL;
1808 }
1809
1810 if (NpyIter_HasDelayedBufAlloc(self->iter)) {
1811 Py_RETURN_TRUE;
1812 }
1813 else {
1814 Py_RETURN_FALSE;
1815 }
1816 }
1817
npyiter_iterationneedsapi_get(NewNpyArrayIterObject * self)1818 static PyObject *npyiter_iterationneedsapi_get(NewNpyArrayIterObject *self)
1819 {
1820 if (self->iter == NULL) {
1821 PyErr_SetString(PyExc_ValueError,
1822 "Iterator is invalid");
1823 return NULL;
1824 }
1825
1826 if (NpyIter_IterationNeedsAPI(self->iter)) {
1827 Py_RETURN_TRUE;
1828 }
1829 else {
1830 Py_RETURN_FALSE;
1831 }
1832 }
1833
npyiter_has_multi_index_get(NewNpyArrayIterObject * self)1834 static PyObject *npyiter_has_multi_index_get(NewNpyArrayIterObject *self)
1835 {
1836 if (self->iter == NULL) {
1837 PyErr_SetString(PyExc_ValueError,
1838 "Iterator is invalid");
1839 return NULL;
1840 }
1841
1842 if (NpyIter_HasMultiIndex(self->iter)) {
1843 Py_RETURN_TRUE;
1844 }
1845 else {
1846 Py_RETURN_FALSE;
1847 }
1848 }
1849
npyiter_has_index_get(NewNpyArrayIterObject * self)1850 static PyObject *npyiter_has_index_get(NewNpyArrayIterObject *self)
1851 {
1852 if (self->iter == NULL) {
1853 PyErr_SetString(PyExc_ValueError,
1854 "Iterator is invalid");
1855 return NULL;
1856 }
1857
1858 if (NpyIter_HasIndex(self->iter)) {
1859 Py_RETURN_TRUE;
1860 }
1861 else {
1862 Py_RETURN_FALSE;
1863 }
1864 }
1865
npyiter_dtypes_get(NewNpyArrayIterObject * self)1866 static PyObject *npyiter_dtypes_get(NewNpyArrayIterObject *self)
1867 {
1868 PyObject *ret;
1869
1870 npy_intp iop, nop;
1871 PyArray_Descr **dtypes;
1872
1873 if (self->iter == NULL) {
1874 PyErr_SetString(PyExc_ValueError,
1875 "Iterator is invalid");
1876 return NULL;
1877 }
1878 nop = NpyIter_GetNOp(self->iter);
1879
1880 ret = PyTuple_New(nop);
1881 if (ret == NULL) {
1882 return NULL;
1883 }
1884 dtypes = self->dtypes;
1885 for (iop = 0; iop < nop; ++iop) {
1886 PyArray_Descr *dtype = dtypes[iop];
1887
1888 Py_INCREF(dtype);
1889 PyTuple_SET_ITEM(ret, iop, (PyObject *)dtype);
1890 }
1891
1892 return ret;
1893 }
1894
npyiter_ndim_get(NewNpyArrayIterObject * self)1895 static PyObject *npyiter_ndim_get(NewNpyArrayIterObject *self)
1896 {
1897 if (self->iter == NULL) {
1898 PyErr_SetString(PyExc_ValueError,
1899 "Iterator is invalid");
1900 return NULL;
1901 }
1902
1903 return PyLong_FromLong(NpyIter_GetNDim(self->iter));
1904 }
1905
npyiter_nop_get(NewNpyArrayIterObject * self)1906 static PyObject *npyiter_nop_get(NewNpyArrayIterObject *self)
1907 {
1908 if (self->iter == NULL) {
1909 PyErr_SetString(PyExc_ValueError,
1910 "Iterator is invalid");
1911 return NULL;
1912 }
1913
1914 return PyLong_FromLong(NpyIter_GetNOp(self->iter));
1915 }
1916
npyiter_itersize_get(NewNpyArrayIterObject * self)1917 static PyObject *npyiter_itersize_get(NewNpyArrayIterObject *self)
1918 {
1919 if (self->iter == NULL) {
1920 PyErr_SetString(PyExc_ValueError,
1921 "Iterator is invalid");
1922 return NULL;
1923 }
1924
1925 return PyLong_FromLong(NpyIter_GetIterSize(self->iter));
1926 }
1927
npyiter_finished_get(NewNpyArrayIterObject * self)1928 static PyObject *npyiter_finished_get(NewNpyArrayIterObject *self)
1929 {
1930 if (self->iter == NULL || !self->finished) {
1931 Py_RETURN_FALSE;
1932 }
1933 else {
1934 Py_RETURN_TRUE;
1935 }
1936 }
1937
1938 NPY_NO_EXPORT Py_ssize_t
npyiter_seq_length(NewNpyArrayIterObject * self)1939 npyiter_seq_length(NewNpyArrayIterObject *self)
1940 {
1941 if (self->iter == NULL) {
1942 return 0;
1943 }
1944 else {
1945 return NpyIter_GetNOp(self->iter);
1946 }
1947 }
1948
1949 NPY_NO_EXPORT PyObject *
npyiter_seq_item(NewNpyArrayIterObject * self,Py_ssize_t i)1950 npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
1951 {
1952 npy_intp ret_ndim;
1953 npy_intp nop, innerloopsize, innerstride;
1954 char *dataptr;
1955 PyArray_Descr *dtype;
1956 int has_external_loop;
1957 Py_ssize_t i_orig = i;
1958
1959 if (self->iter == NULL || self->finished) {
1960 PyErr_SetString(PyExc_ValueError,
1961 "Iterator is past the end");
1962 return NULL;
1963 }
1964
1965 if (NpyIter_HasDelayedBufAlloc(self->iter)) {
1966 PyErr_SetString(PyExc_ValueError,
1967 "Iterator construction used delayed buffer allocation, "
1968 "and no reset has been done yet");
1969 return NULL;
1970 }
1971 nop = NpyIter_GetNOp(self->iter);
1972
1973 /* Negative indexing */
1974 if (i < 0) {
1975 i += nop;
1976 }
1977
1978 if (i < 0 || i >= nop) {
1979 PyErr_Format(PyExc_IndexError,
1980 "Iterator operand index %zd is out of bounds", i_orig);
1981 return NULL;
1982 }
1983
1984 #if 0
1985 /*
1986 * This check is disabled because it prevents things like
1987 * np.add(it[0], it[1], it[2]), where it[2] is a write-only
1988 * parameter. When write-only, the value of it[i] is
1989 * likely random junk, as if it were allocated with an
1990 * np.empty(...) call.
1991 */
1992 if (!self->readflags[i]) {
1993 PyErr_Format(PyExc_RuntimeError,
1994 "Iterator operand %zd is write-only", i);
1995 return NULL;
1996 }
1997 #endif
1998
1999 dataptr = self->dataptrs[i];
2000 dtype = self->dtypes[i];
2001 has_external_loop = NpyIter_HasExternalLoop(self->iter);
2002
2003 if (has_external_loop) {
2004 innerloopsize = *self->innerloopsizeptr;
2005 innerstride = self->innerstrides[i];
2006 ret_ndim = 1;
2007 }
2008 else {
2009 innerloopsize = 1;
2010 innerstride = 0;
2011 /* If the iterator is going over every element, return array scalars */
2012 ret_ndim = 0;
2013 }
2014
2015 Py_INCREF(dtype);
2016 return PyArray_NewFromDescrAndBase(
2017 &PyArray_Type, dtype,
2018 ret_ndim, &innerloopsize, &innerstride, dataptr,
2019 self->writeflags[i] ? NPY_ARRAY_WRITEABLE : 0,
2020 NULL, (PyObject *)self);
2021 }
2022
2023 NPY_NO_EXPORT PyObject *
npyiter_seq_slice(NewNpyArrayIterObject * self,Py_ssize_t ilow,Py_ssize_t ihigh)2024 npyiter_seq_slice(NewNpyArrayIterObject *self,
2025 Py_ssize_t ilow, Py_ssize_t ihigh)
2026 {
2027 PyObject *ret;
2028 npy_intp nop;
2029 Py_ssize_t i;
2030
2031 if (self->iter == NULL || self->finished) {
2032 PyErr_SetString(PyExc_ValueError,
2033 "Iterator is past the end");
2034 return NULL;
2035 }
2036
2037 if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2038 PyErr_SetString(PyExc_ValueError,
2039 "Iterator construction used delayed buffer allocation, "
2040 "and no reset has been done yet");
2041 return NULL;
2042 }
2043 nop = NpyIter_GetNOp(self->iter);
2044 if (ilow < 0) {
2045 ilow = 0;
2046 }
2047 else if (ilow >= nop) {
2048 ilow = nop-1;
2049 }
2050 if (ihigh < ilow) {
2051 ihigh = ilow;
2052 }
2053 else if (ihigh > nop) {
2054 ihigh = nop;
2055 }
2056
2057 ret = PyTuple_New(ihigh-ilow);
2058 if (ret == NULL) {
2059 return NULL;
2060 }
2061 for (i = ilow; i < ihigh ; ++i) {
2062 PyObject *item = npyiter_seq_item(self, i);
2063 if (item == NULL) {
2064 Py_DECREF(ret);
2065 return NULL;
2066 }
2067 PyTuple_SET_ITEM(ret, i-ilow, item);
2068 }
2069 return ret;
2070 }
2071
2072 NPY_NO_EXPORT int
npyiter_seq_ass_item(NewNpyArrayIterObject * self,Py_ssize_t i,PyObject * v)2073 npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
2074 {
2075
2076 npy_intp nop, innerloopsize, innerstride;
2077 char *dataptr;
2078 PyArray_Descr *dtype;
2079 PyArrayObject *tmp;
2080 int ret, has_external_loop;
2081 Py_ssize_t i_orig = i;
2082
2083
2084 if (v == NULL) {
2085 PyErr_SetString(PyExc_TypeError,
2086 "Cannot delete iterator elements");
2087 return -1;
2088 }
2089
2090 if (self->iter == NULL || self->finished) {
2091 PyErr_SetString(PyExc_ValueError,
2092 "Iterator is past the end");
2093 return -1;
2094 }
2095
2096 if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2097 PyErr_SetString(PyExc_ValueError,
2098 "Iterator construction used delayed buffer allocation, "
2099 "and no reset has been done yet");
2100 return -1;
2101 }
2102 nop = NpyIter_GetNOp(self->iter);
2103
2104 /* Negative indexing */
2105 if (i < 0) {
2106 i += nop;
2107 }
2108
2109 if (i < 0 || i >= nop) {
2110 PyErr_Format(PyExc_IndexError,
2111 "Iterator operand index %zd is out of bounds", i_orig);
2112 return -1;
2113 }
2114 if (!self->writeflags[i]) {
2115 PyErr_Format(PyExc_RuntimeError,
2116 "Iterator operand %zd is not writeable", i_orig);
2117 return -1;
2118 }
2119
2120 dataptr = self->dataptrs[i];
2121 dtype = self->dtypes[i];
2122 has_external_loop = NpyIter_HasExternalLoop(self->iter);
2123
2124 if (has_external_loop) {
2125 innerloopsize = *self->innerloopsizeptr;
2126 innerstride = self->innerstrides[i];
2127 }
2128 else {
2129 innerloopsize = 1;
2130 innerstride = 0;
2131 }
2132
2133 /* TODO - there should be a better way than this... */
2134 Py_INCREF(dtype);
2135 tmp = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype,
2136 1, &innerloopsize,
2137 &innerstride, dataptr,
2138 NPY_ARRAY_WRITEABLE, NULL);
2139 if (tmp == NULL) {
2140 return -1;
2141 }
2142
2143 ret = PyArray_CopyObject(tmp, v);
2144 Py_DECREF(tmp);
2145 return ret;
2146 }
2147
2148 static int
npyiter_seq_ass_slice(NewNpyArrayIterObject * self,Py_ssize_t ilow,Py_ssize_t ihigh,PyObject * v)2149 npyiter_seq_ass_slice(NewNpyArrayIterObject *self, Py_ssize_t ilow,
2150 Py_ssize_t ihigh, PyObject *v)
2151 {
2152 npy_intp nop;
2153 Py_ssize_t i;
2154
2155 if (v == NULL) {
2156 PyErr_SetString(PyExc_TypeError,
2157 "Cannot delete iterator elements");
2158 return -1;
2159 }
2160
2161 if (self->iter == NULL || self->finished) {
2162 PyErr_SetString(PyExc_ValueError,
2163 "Iterator is past the end");
2164 return -1;
2165 }
2166
2167 if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2168 PyErr_SetString(PyExc_ValueError,
2169 "Iterator construction used delayed buffer allocation, "
2170 "and no reset has been done yet");
2171 return -1;
2172 }
2173 nop = NpyIter_GetNOp(self->iter);
2174 if (ilow < 0) {
2175 ilow = 0;
2176 }
2177 else if (ilow >= nop) {
2178 ilow = nop-1;
2179 }
2180 if (ihigh < ilow) {
2181 ihigh = ilow;
2182 }
2183 else if (ihigh > nop) {
2184 ihigh = nop;
2185 }
2186
2187 if (!PySequence_Check(v) || PySequence_Size(v) != ihigh-ilow) {
2188 PyErr_SetString(PyExc_ValueError,
2189 "Wrong size to assign to iterator slice");
2190 return -1;
2191 }
2192
2193 for (i = ilow; i < ihigh ; ++i) {
2194 PyObject *item = PySequence_GetItem(v, i-ilow);
2195 if (item == NULL) {
2196 return -1;
2197 }
2198 if (npyiter_seq_ass_item(self, i, item) < 0) {
2199 Py_DECREF(item);
2200 return -1;
2201 }
2202 Py_DECREF(item);
2203 }
2204
2205 return 0;
2206 }
2207
2208 static PyObject *
npyiter_subscript(NewNpyArrayIterObject * self,PyObject * op)2209 npyiter_subscript(NewNpyArrayIterObject *self, PyObject *op)
2210 {
2211 if (self->iter == NULL || self->finished) {
2212 PyErr_SetString(PyExc_ValueError,
2213 "Iterator is past the end");
2214 return NULL;
2215 }
2216
2217 if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2218 PyErr_SetString(PyExc_ValueError,
2219 "Iterator construction used delayed buffer allocation, "
2220 "and no reset has been done yet");
2221 return NULL;
2222 }
2223
2224 if (PyLong_Check(op) ||
2225 (PyIndex_Check(op) && !PySequence_Check(op))) {
2226 npy_intp i = PyArray_PyIntAsIntp(op);
2227 if (error_converting(i)) {
2228 return NULL;
2229 }
2230 return npyiter_seq_item(self, i);
2231 }
2232 else if (PySlice_Check(op)) {
2233 Py_ssize_t istart = 0, iend = 0, istep = 0, islicelength;
2234 if (PySlice_GetIndicesEx(op, NpyIter_GetNOp(self->iter),
2235 &istart, &iend, &istep, &islicelength) < 0) {
2236 return NULL;
2237 }
2238 if (istep != 1) {
2239 PyErr_SetString(PyExc_ValueError,
2240 "Iterator slicing only supports a step of 1");
2241 return NULL;
2242 }
2243 return npyiter_seq_slice(self, istart, iend);
2244 }
2245
2246 PyErr_SetString(PyExc_TypeError,
2247 "invalid index type for iterator indexing");
2248 return NULL;
2249 }
2250
2251 static int
npyiter_ass_subscript(NewNpyArrayIterObject * self,PyObject * op,PyObject * value)2252 npyiter_ass_subscript(NewNpyArrayIterObject *self, PyObject *op,
2253 PyObject *value)
2254 {
2255 if (value == NULL) {
2256 PyErr_SetString(PyExc_TypeError,
2257 "Cannot delete iterator elements");
2258 return -1;
2259 }
2260 if (self->iter == NULL || self->finished) {
2261 PyErr_SetString(PyExc_ValueError,
2262 "Iterator is past the end");
2263 return -1;
2264 }
2265
2266 if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2267 PyErr_SetString(PyExc_ValueError,
2268 "Iterator construction used delayed buffer allocation, "
2269 "and no reset has been done yet");
2270 return -1;
2271 }
2272
2273 if (PyLong_Check(op) ||
2274 (PyIndex_Check(op) && !PySequence_Check(op))) {
2275 npy_intp i = PyArray_PyIntAsIntp(op);
2276 if (error_converting(i)) {
2277 return -1;
2278 }
2279 return npyiter_seq_ass_item(self, i, value);
2280 }
2281 else if (PySlice_Check(op)) {
2282 Py_ssize_t istart = 0, iend = 0, istep = 0, islicelength = 0;
2283 if (PySlice_GetIndicesEx(op, NpyIter_GetNOp(self->iter),
2284 &istart, &iend, &istep, &islicelength) < 0) {
2285 return -1;
2286 }
2287 if (istep != 1) {
2288 PyErr_SetString(PyExc_ValueError,
2289 "Iterator slice assignment only supports a step of 1");
2290 return -1;
2291 }
2292 return npyiter_seq_ass_slice(self, istart, iend, value);
2293 }
2294
2295 PyErr_SetString(PyExc_TypeError,
2296 "invalid index type for iterator indexing");
2297 return -1;
2298 }
2299
2300 static PyObject *
npyiter_enter(NewNpyArrayIterObject * self)2301 npyiter_enter(NewNpyArrayIterObject *self)
2302 {
2303 if (self->iter == NULL) {
2304 PyErr_SetString(PyExc_RuntimeError, "operation on non-initialized iterator");
2305 return NULL;
2306 }
2307 Py_INCREF(self);
2308 return (PyObject *)self;
2309 }
2310
2311 static PyObject *
npyiter_close(NewNpyArrayIterObject * self)2312 npyiter_close(NewNpyArrayIterObject *self)
2313 {
2314 NpyIter *iter = self->iter;
2315 int ret;
2316 if (self->iter == NULL) {
2317 Py_RETURN_NONE;
2318 }
2319 ret = NpyIter_Deallocate(iter);
2320 self->iter = NULL;
2321 Py_XDECREF(self->nested_child);
2322 self->nested_child = NULL;
2323 if (ret < 0) {
2324 return NULL;
2325 }
2326 Py_RETURN_NONE;
2327 }
2328
2329 static PyObject *
npyiter_exit(NewNpyArrayIterObject * self,PyObject * NPY_UNUSED (args))2330 npyiter_exit(NewNpyArrayIterObject *self, PyObject *NPY_UNUSED(args))
2331 {
2332 /* even if called via exception handling, writeback any data */
2333 return npyiter_close(self);
2334 }
2335
2336 static PyMethodDef npyiter_methods[] = {
2337 {"reset",
2338 (PyCFunction)npyiter_reset,
2339 METH_NOARGS, NULL},
2340 {"copy",
2341 (PyCFunction)npyiter_copy,
2342 METH_NOARGS, NULL},
2343 {"__copy__",
2344 (PyCFunction)npyiter_copy,
2345 METH_NOARGS, NULL},
2346 {"iternext",
2347 (PyCFunction)npyiter_iternext,
2348 METH_NOARGS, NULL},
2349 {"remove_axis",
2350 (PyCFunction)npyiter_remove_axis,
2351 METH_VARARGS, NULL},
2352 {"remove_multi_index",
2353 (PyCFunction)npyiter_remove_multi_index,
2354 METH_NOARGS, NULL},
2355 {"enable_external_loop",
2356 (PyCFunction)npyiter_enable_external_loop,
2357 METH_NOARGS, NULL},
2358 {"debug_print",
2359 (PyCFunction)npyiter_debug_print,
2360 METH_NOARGS, NULL},
2361 {"__enter__", (PyCFunction)npyiter_enter,
2362 METH_NOARGS, NULL},
2363 {"__exit__", (PyCFunction)npyiter_exit,
2364 METH_VARARGS, NULL},
2365 {"close", (PyCFunction)npyiter_close,
2366 METH_VARARGS, NULL},
2367 {NULL, NULL, 0, NULL},
2368 };
2369
2370 static PyMemberDef npyiter_members[] = {
2371 {NULL, 0, 0, 0, NULL},
2372 };
2373
2374 static PyGetSetDef npyiter_getsets[] = {
2375 {"value",
2376 (getter)npyiter_value_get,
2377 NULL, NULL, NULL},
2378 {"shape",
2379 (getter)npyiter_shape_get,
2380 NULL, NULL, NULL},
2381 {"multi_index",
2382 (getter)npyiter_multi_index_get,
2383 (setter)npyiter_multi_index_set,
2384 NULL, NULL},
2385 {"index",
2386 (getter)npyiter_index_get,
2387 (setter)npyiter_index_set,
2388 NULL, NULL},
2389 {"iterindex",
2390 (getter)npyiter_iterindex_get,
2391 (setter)npyiter_iterindex_set,
2392 NULL, NULL},
2393 {"iterrange",
2394 (getter)npyiter_iterrange_get,
2395 (setter)npyiter_iterrange_set,
2396 NULL, NULL},
2397 {"operands",
2398 (getter)npyiter_operands_get,
2399 NULL, NULL, NULL},
2400 {"itviews",
2401 (getter)npyiter_itviews_get,
2402 NULL, NULL, NULL},
2403 {"has_delayed_bufalloc",
2404 (getter)npyiter_has_delayed_bufalloc_get,
2405 NULL, NULL, NULL},
2406 {"iterationneedsapi",
2407 (getter)npyiter_iterationneedsapi_get,
2408 NULL, NULL, NULL},
2409 {"has_multi_index",
2410 (getter)npyiter_has_multi_index_get,
2411 NULL, NULL, NULL},
2412 {"has_index",
2413 (getter)npyiter_has_index_get,
2414 NULL, NULL, NULL},
2415 {"dtypes",
2416 (getter)npyiter_dtypes_get,
2417 NULL, NULL, NULL},
2418 {"ndim",
2419 (getter)npyiter_ndim_get,
2420 NULL, NULL, NULL},
2421 {"nop",
2422 (getter)npyiter_nop_get,
2423 NULL, NULL, NULL},
2424 {"itersize",
2425 (getter)npyiter_itersize_get,
2426 NULL, NULL, NULL},
2427 {"finished",
2428 (getter)npyiter_finished_get,
2429 NULL, NULL, NULL},
2430
2431 {NULL, NULL, NULL, NULL, NULL}
2432 };
2433
2434 NPY_NO_EXPORT PySequenceMethods npyiter_as_sequence = {
2435 (lenfunc)npyiter_seq_length, /*sq_length*/
2436 (binaryfunc)NULL, /*sq_concat*/
2437 (ssizeargfunc)NULL, /*sq_repeat*/
2438 (ssizeargfunc)npyiter_seq_item, /*sq_item*/
2439 (ssizessizeargfunc)NULL, /*sq_slice*/
2440 (ssizeobjargproc)npyiter_seq_ass_item, /*sq_ass_item*/
2441 (ssizessizeobjargproc)NULL, /*sq_ass_slice*/
2442 (objobjproc)NULL, /*sq_contains */
2443 (binaryfunc)NULL, /*sq_inplace_concat */
2444 (ssizeargfunc)NULL, /*sq_inplace_repeat */
2445 };
2446
2447 NPY_NO_EXPORT PyMappingMethods npyiter_as_mapping = {
2448 (lenfunc)npyiter_seq_length, /*mp_length*/
2449 (binaryfunc)npyiter_subscript, /*mp_subscript*/
2450 (objobjargproc)npyiter_ass_subscript, /*mp_ass_subscript*/
2451 };
2452
2453 NPY_NO_EXPORT PyTypeObject NpyIter_Type = {
2454 PyVarObject_HEAD_INIT(NULL, 0)
2455 .tp_name = "numpy.nditer",
2456 .tp_basicsize = sizeof(NewNpyArrayIterObject),
2457 .tp_dealloc = (destructor)npyiter_dealloc,
2458 .tp_as_sequence = &npyiter_as_sequence,
2459 .tp_as_mapping = &npyiter_as_mapping,
2460 .tp_flags = Py_TPFLAGS_DEFAULT,
2461 .tp_iternext = (iternextfunc)npyiter_next,
2462 .tp_methods = npyiter_methods,
2463 .tp_members = npyiter_members,
2464 .tp_getset = npyiter_getsets,
2465 .tp_init = (initproc)npyiter_init,
2466 .tp_new = npyiter_new,
2467 };
2468