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

48 1
static int npyiter_cache_values(NewNpyArrayIterObject *self)
49
{
50 1
    NpyIter *iter = self->iter;
51

52
    /* iternext and get_multi_index functions */
53 1
    self->iternext = NpyIter_GetIterNext(iter, NULL);
54 1
    if (self->iternext == NULL) {
55
        return -1;
56
    }
57

58 1
    if (NpyIter_HasMultiIndex(iter) && !NpyIter_HasDelayedBufAlloc(iter)) {
59 1
        self->get_multi_index = NpyIter_GetGetMultiIndex(iter, NULL);
60
    }
61
    else {
62 1
        self->get_multi_index = NULL;
63
    }
64

65
    /* Internal data pointers */
66 1
    self->dataptrs = NpyIter_GetDataPtrArray(iter);
67 1
    self->dtypes = NpyIter_GetDescrArray(iter);
68 1
    self->operands = NpyIter_GetOperandArray(iter);
69

70 1
    if (NpyIter_HasExternalLoop(iter)) {
71 1
        self->innerstrides = NpyIter_GetInnerStrideArray(iter);
72 1
        self->innerloopsizeptr = NpyIter_GetInnerLoopSizePtr(iter);
73
    }
74
    else {
75 1
        self->innerstrides = NULL;
76 1
        self->innerloopsizeptr = NULL;
77
    }
78

79
    /* The read/write settings */
80 1
    NpyIter_GetReadFlags(iter, self->readflags);
81 1
    NpyIter_GetWriteFlags(iter, self->writeflags);
82 1
    return 0;
83
}
84

85
static PyObject *
86 1
npyiter_new(PyTypeObject *subtype, PyObject *NPY_UNUSED(args),
87
            PyObject *NPY_UNUSED(kwds))
88
{
89
    NewNpyArrayIterObject *self;
90

91 1
    self = (NewNpyArrayIterObject *)subtype->tp_alloc(subtype, 0);
92 1
    if (self != NULL) {
93 1
        self->iter = NULL;
94 1
        self->nested_child = NULL;
95
    }
96

97 1
    return (PyObject *)self;
98
}
99

100
static int
101 1
NpyIter_GlobalFlagsConverter(PyObject *flags_in, npy_uint32 *flags)
102
{
103 1
    npy_uint32 tmpflags = 0;
104
    int iflags, nflags;
105

106
    PyObject *f;
107 1
    char *str = NULL;
108 1
    Py_ssize_t length = 0;
109
    npy_uint32 flag;
110

111 1
    if (flags_in == NULL || flags_in == Py_None) {
112
        return 1;
113
    }
114

115 1
    if (!PyTuple_Check(flags_in) && !PyList_Check(flags_in)) {
116 0
        PyErr_SetString(PyExc_ValueError,
117
                "Iterator global flags must be a list or tuple of strings");
118 0
        return 0;
119
    }
120

121 1
    nflags = PySequence_Size(flags_in);
122

123 1
    for (iflags = 0; iflags < nflags; ++iflags) {
124 1
        f = PySequence_GetItem(flags_in, iflags);
125 1
        if (f == NULL) {
126
            return 0;
127
        }
128

129 1
        if (PyUnicode_Check(f)) {
130
            /* accept unicode input */
131
            PyObject *f_str;
132 1
            f_str = PyUnicode_AsASCIIString(f);
133 1
            if (f_str == NULL) {
134 0
                Py_DECREF(f);
135
                return 0;
136
            }
137 1
            Py_DECREF(f);
138
            f = f_str;
139
        }
140

141 1
        if (PyBytes_AsStringAndSize(f, &str, &length) < 0) {
142 0
            Py_DECREF(f);
143
            return 0;
144
        }
145
        /* Use switch statements to quickly isolate the right flag */
146 1
        flag = 0;
147 1
        switch (str[0]) {
148 1
            case 'b':
149 1
                if (strcmp(str, "buffered") == 0) {
150
                    flag = NPY_ITER_BUFFERED;
151
                }
152
                break;
153 1
            case 'c':
154 1
                if (length >= 6) switch (str[5]) {
155 1
                    case 'e':
156 1
                        if (strcmp(str, "c_index") == 0) {
157
                            flag = NPY_ITER_C_INDEX;
158
                        }
159
                        break;
160 1
                    case 'i':
161 1
                        if (strcmp(str, "copy_if_overlap") == 0) {
162
                            flag = NPY_ITER_COPY_IF_OVERLAP;
163
                        }
164
                        break;
165 1
                    case 'n':
166 1
                        if (strcmp(str, "common_dtype") == 0) {
167
                            flag = NPY_ITER_COMMON_DTYPE;
168
                        }
169
                        break;
170
                }
171
                break;
172 1
            case 'd':
173 1
                if (strcmp(str, "delay_bufalloc") == 0) {
174
                    flag = NPY_ITER_DELAY_BUFALLOC;
175
                }
176
                break;
177 1
            case 'e':
178 1
                if (strcmp(str, "external_loop") == 0) {
179
                    flag = NPY_ITER_EXTERNAL_LOOP;
180
                }
181
                break;
182 1
            case 'f':
183 1
                if (strcmp(str, "f_index") == 0) {
184
                    flag = NPY_ITER_F_INDEX;
185
                }
186
                break;
187 1
            case 'g':
188
                /*
189
                 * Documentation is grow_inner, but initial implementation
190
                 * was growinner, so allowing for either.
191
                 */
192 1
                if (strcmp(str, "grow_inner") == 0 ||
193 1
                            strcmp(str, "growinner") == 0) {
194
                    flag = NPY_ITER_GROWINNER;
195
                }
196
                break;
197 1
            case 'm':
198 1
                if (strcmp(str, "multi_index") == 0) {
199
                    flag = NPY_ITER_MULTI_INDEX;
200
                }
201
                break;
202 1
            case 'r':
203 1
                if (strcmp(str, "ranged") == 0) {
204
                    flag = NPY_ITER_RANGED;
205
                }
206 1
                else if (strcmp(str, "refs_ok") == 0) {
207
                    flag = NPY_ITER_REFS_OK;
208
                }
209 1
                else if (strcmp(str, "reduce_ok") == 0) {
210
                    flag = NPY_ITER_REDUCE_OK;
211
                }
212
                break;
213 1
            case 'z':
214 1
                if (strcmp(str, "zerosize_ok") == 0) {
215
                    flag = NPY_ITER_ZEROSIZE_OK;
216
                }
217
                break;
218
        }
219
        if (flag == 0) {
220 1
            PyErr_Format(PyExc_ValueError,
221
                    "Unexpected iterator global flag \"%s\"", str);
222 1
            Py_DECREF(f);
223
            return 0;
224
        }
225
        else {
226 1
            tmpflags |= flag;
227
        }
228 1
        Py_DECREF(f);
229
    }
230

231 1
    *flags |= tmpflags;
232 1
    return 1;
233
}
234

235
static int
236 1
NpyIter_OpFlagsConverter(PyObject *op_flags_in,
237
                         npy_uint32 *op_flags)
238
{
239
    int iflags, nflags;
240
    npy_uint32 flag;
241

242 1
    if (!PyTuple_Check(op_flags_in) && !PyList_Check(op_flags_in)) {
243 0
        PyErr_SetString(PyExc_ValueError,
244
                "op_flags must be a tuple or array of per-op flag-tuples");
245 0
        return 0;
246
    }
247

248 1
    nflags = PySequence_Size(op_flags_in);
249

250 1
    *op_flags = 0;
251 1
    for (iflags = 0; iflags < nflags; ++iflags) {
252
        PyObject *f;
253 1
        char *str = NULL;
254 1
        Py_ssize_t length = 0;
255

256 1
        f = PySequence_GetItem(op_flags_in, iflags);
257 1
        if (f == NULL) {
258 1
            return 0;
259
        }
260

261 1
        if (PyUnicode_Check(f)) {
262
            /* accept unicode input */
263
            PyObject *f_str;
264 1
            f_str = PyUnicode_AsASCIIString(f);
265 1
            if (f_str == NULL) {
266 0
                Py_DECREF(f);
267
                return 0;
268
            }
269 1
            Py_DECREF(f);
270
            f = f_str;
271
        }
272

273 1
        if (PyBytes_AsStringAndSize(f, &str, &length) < 0) {
274 1
            PyErr_Clear();
275 1
            Py_DECREF(f);
276 1
            PyErr_SetString(PyExc_ValueError,
277
                   "op_flags must be a tuple or array of per-op flag-tuples");
278 1
            return 0;
279
        }
280

281
        /* Use switch statements to quickly isolate the right flag */
282 1
        flag = 0;
283 1
        switch (str[0]) {
284 1
            case 'a':
285 1
                if (length > 2) switch(str[2]) {
286 1
                    case 'i':
287 1
                        if (strcmp(str, "aligned") == 0) {
288
                            flag = NPY_ITER_ALIGNED;
289
                        }
290
                        break;
291 1
                    case 'l':
292 1
                        if (strcmp(str, "allocate") == 0) {
293
                            flag = NPY_ITER_ALLOCATE;
294
                        }
295
                        break;
296 1
                    case 'r':
297 1
                        if (strcmp(str, "arraymask") == 0) {
298
                            flag = NPY_ITER_ARRAYMASK;
299
                        }
300
                        break;
301
                }
302
                break;
303 1
            case 'c':
304 1
                if (strcmp(str, "copy") == 0) {
305 1
                    flag = NPY_ITER_COPY;
306
                }
307 1
                if (strcmp(str, "contig") == 0) {
308
                    flag = NPY_ITER_CONTIG;
309
                }
310
                break;
311 1
            case 'n':
312 1
                switch (str[1]) {
313 1
                    case 'b':
314 1
                        if (strcmp(str, "nbo") == 0) {
315
                            flag = NPY_ITER_NBO;
316
                        }
317
                        break;
318 1
                    case 'o':
319 1
                        if (strcmp(str, "no_subtype") == 0) {
320
                            flag = NPY_ITER_NO_SUBTYPE;
321
                        }
322 1
                        else if (strcmp(str, "no_broadcast") == 0) {
323
                            flag = NPY_ITER_NO_BROADCAST;
324
                        }
325
                        break;
326
                }
327
                break;
328 1
            case 'o':
329 1
                if (strcmp(str, "overlap_assume_elementwise") == 0) {
330
                    flag = NPY_ITER_OVERLAP_ASSUME_ELEMENTWISE;
331
                }
332
                break;
333 1
            case 'r':
334 1
                if (length > 4) switch (str[4]) {
335 1
                    case 'o':
336 1
                        if (strcmp(str, "readonly") == 0) {
337
                            flag = NPY_ITER_READONLY;
338
                        }
339
                        break;
340 1
                    case 'w':
341 1
                        if (strcmp(str, "readwrite") == 0) {
342
                            flag = NPY_ITER_READWRITE;
343
                        }
344
                        break;
345
                }
346
                break;
347 1
            case 'u':
348 1
                switch (str[1]) {
349 1
                    case 'p':
350 1
                        if (strcmp(str, "updateifcopy") == 0) {
351
                            flag = NPY_ITER_UPDATEIFCOPY;
352
                        }
353
                        break;
354
                }
355
                break;
356 0
            case 'v':
357 0
                if (strcmp(str, "virtual") == 0) {
358
                    flag = NPY_ITER_VIRTUAL;
359
                }
360
                break;
361 1
            case 'w':
362 1
                if (length > 5) switch (str[5]) {
363 1
                    case 'o':
364 1
                        if (strcmp(str, "writeonly") == 0) {
365
                            flag = NPY_ITER_WRITEONLY;
366
                        }
367
                        break;
368 1
                    case 'm':
369 1
                        if (strcmp(str, "writemasked") == 0) {
370
                            flag = NPY_ITER_WRITEMASKED;
371
                        }
372
                        break;
373
                }
374
                break;
375
        }
376 1
        if (flag == 0) {
377 1
            PyErr_Format(PyExc_ValueError,
378
                    "Unexpected per-op iterator flag \"%s\"", str);
379 1
            Py_DECREF(f);
380
            return 0;
381
        }
382
        else {
383 1
            *op_flags |= flag;
384
        }
385 1
        Py_DECREF(f);
386
    }
387

388
    return 1;
389
}
390

391
static int
392 1
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 1
    if (!PyTuple_Check(op_flags_in) && !PyList_Check(op_flags_in)) {
398 0
        PyErr_SetString(PyExc_ValueError,
399
                "op_flags must be a tuple or array of per-op flag-tuples");
400 0
        return 0;
401
    }
402

403 1
    if (PySequence_Size(op_flags_in) != nop) {
404
        goto try_single_flags;
405
    }
406

407 1
    for (iop = 0; iop < nop; ++iop) {
408 1
        PyObject *f = PySequence_GetItem(op_flags_in, iop);
409 1
        if (f == NULL) {
410
            return 0;
411
        }
412
        /* If the first item is a string, try as one set of flags */
413 1
        if (iop == 0 && (PyBytes_Check(f) || PyUnicode_Check(f))) {
414 1
            Py_DECREF(f);
415
            goto try_single_flags;
416
        }
417 1
        if (NpyIter_OpFlagsConverter(f,
418 1
                        &op_flags_array[iop]) != 1) {
419 1
            Py_DECREF(f);
420
            return 0;
421
        }
422

423 1
        Py_DECREF(f);
424
    }
425

426
    return 1;
427

428 1
try_single_flags:
429 1
    if (NpyIter_OpFlagsConverter(op_flags_in,
430
                        &op_flags_array[0]) != 1) {
431
        return 0;
432
    }
433

434 1
    for (iop = 1; iop < nop; ++iop) {
435 1
        op_flags_array[iop] = op_flags_array[0];
436
    }
437

438
    return 1;
439
}
440

441
static int
442 1
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 1
    if ((!PyTuple_Check(op_dtypes_in) && !PyList_Check(op_dtypes_in)) ||
453 1
                                    PySequence_Size(op_dtypes_in) != nop) {
454
        goto try_single_dtype;
455
    }
456

457 1
    for (iop = 0; iop < nop; ++iop) {
458 1
        PyObject *dtype = PySequence_GetItem(op_dtypes_in, iop);
459 1
        if (dtype == NULL) {
460
            npy_intp i;
461 0
            for (i = 0; i < iop; ++i ) {
462 0
                Py_XDECREF(op_dtypes[i]);
463
            }
464
            return 0;
465
        }
466

467
        /* Try converting the object to a descr */
468 1
        if (PyArray_DescrConverter2(dtype, &op_dtypes[iop]) != 1) {
469
            npy_intp i;
470 0
            for (i = 0; i < iop; ++i ) {
471 0
                Py_XDECREF(op_dtypes[i]);
472
            }
473 1
            Py_DECREF(dtype);
474 1
            PyErr_Clear();
475 1
            goto try_single_dtype;
476
        }
477

478 1
        Py_DECREF(dtype);
479
    }
480

481
    return 1;
482

483 1
try_single_dtype:
484 1
    if (PyArray_DescrConverter2(op_dtypes_in, &op_dtypes[0]) == 1) {
485 1
        for (iop = 1; iop < nop; ++iop) {
486 1
            op_dtypes[iop] = op_dtypes[0];
487 1
            Py_XINCREF(op_dtypes[iop]);
488
        }
489
        return 1;
490
    }
491

492
    return 0;
493
}
494

495
static int
496 1
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 1
    if ((!PyTuple_Check(op_axes_in) && !PyList_Check(op_axes_in)) ||
503 1
                                PySequence_Size(op_axes_in) != nop) {
504 1
        PyErr_SetString(PyExc_ValueError,
505
                "op_axes must be a tuple/list matching the number of ops");
506 1
        return 0;
507
    }
508

509 1
    *oa_ndim = -1;
510

511
    /* Copy the tuples into op_axes */
512 1
    for (iop = 0; iop < nop; ++iop) {
513
        int idim;
514 1
        a = PySequence_GetItem(op_axes_in, iop);
515 1
        if (a == NULL) {
516
            return 0;
517
        }
518 1
        if (a == Py_None) {
519 1
            op_axes[iop] = NULL;
520
        } else {
521 1
            if (!PyTuple_Check(a) && !PyList_Check(a)) {
522 0
                PyErr_SetString(PyExc_ValueError,
523
                        "Each entry of op_axes must be None "
524
                        "or a tuple/list");
525 0
                Py_DECREF(a);
526
                return 0;
527
            }
528 1
            if (*oa_ndim == -1) {
529 1
                *oa_ndim = PySequence_Size(a);
530 1
                if (*oa_ndim > NPY_MAXDIMS) {
531 0
                    PyErr_SetString(PyExc_ValueError,
532
                            "Too many dimensions in op_axes");
533 0
                    Py_DECREF(a);
534
                    return 0;
535
                }
536
            }
537 1
            if (PySequence_Size(a) != *oa_ndim) {
538 1
                PyErr_SetString(PyExc_ValueError,
539
                        "Each entry of op_axes must have the same size");
540 1
                Py_DECREF(a);
541
                return 0;
542
            }
543 1
            for (idim = 0; idim < *oa_ndim; ++idim) {
544 1
                PyObject *v = PySequence_GetItem(a, idim);
545 1
                if (v == NULL) {
546 0
                    Py_DECREF(a);
547
                    return 0;
548
                }
549
                /* numpy.newaxis is None */
550 1
                if (v == Py_None) {
551 1
                    op_axes[iop][idim] = -1;
552
                }
553
                else {
554 1
                    op_axes[iop][idim] = PyArray_PyIntAsInt(v);
555 1
                    if (op_axes[iop][idim]==-1 &&
556 1
                                                PyErr_Occurred()) {
557 0
                        Py_DECREF(a);
558 0
                        Py_DECREF(v);
559
                        return 0;
560
                    }
561
                }
562 1
                Py_DECREF(v);
563
            }
564
        }
565 1
        Py_DECREF(a);
566
    }
567

568 1
    if (*oa_ndim == -1) {
569 0
        PyErr_SetString(PyExc_ValueError,
570
                "If op_axes is provided, at least one list of axes "
571
                "must be contained within it");
572 0
        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
584 1
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 1
    if (PyTuple_Check(op_in) || PyList_Check(op_in)) {
592 1
        nop = PySequence_Size(op_in);
593 1
        if (nop == 0) {
594 1
            PyErr_SetString(PyExc_ValueError,
595
                    "Must provide at least one operand");
596 1
            return 0;
597
        }
598 1
        if (nop > NPY_MAXARGS) {
599 1
            PyErr_SetString(PyExc_ValueError, "Too many operands");
600 1
            return 0;
601
        }
602

603 1
        for (iop = 0; iop < nop; ++iop) {
604 1
            PyObject *item = PySequence_GetItem(op_in, iop);
605 1
            if (item == NULL) {
606
                npy_intp i;
607 0
                for (i = 0; i < iop; ++i) {
608 0
                    Py_XDECREF(op[i]);
609
                }
610
                return 0;
611
            }
612 1
            else if (item == Py_None) {
613 1
                Py_DECREF(item);
614
                item = NULL;
615
            }
616
            /* This is converted to an array after op flags are retrieved */
617 1
            op[iop] = (PyArrayObject *)item;
618
        }
619
    }
620
    else {
621 1
        nop = 1;
622
        /* Is converted to an array after op flags are retrieved */
623 1
        Py_INCREF(op_in);
624 1
        op[0] = (PyArrayObject *)op_in;
625
    }
626

627 1
    *nop_out = nop;
628

629
    /* op_flags */
630 1
    if (op_flags_in == NULL || op_flags_in == Py_None) {
631 1
        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 1
            if (op[iop] == NULL) {
638 1
                op_flags[iop] = NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE;
639
            }
640
            else {
641 1
                op_flags[iop] = NPY_ITER_READONLY;
642
            }
643
        }
644
    }
645 1
    else if (npyiter_convert_op_flags_array(op_flags_in,
646
                                      op_flags, nop) != 1) {
647 1
        for (iop = 0; iop < nop; ++iop) {
648 1
            Py_XDECREF(op[iop]);
649
        }
650 1
        *nop_out = 0;
651 1
        return 0;
652
    }
653

654
    /* Now that we have the flags - convert all the ops to arrays */
655 1
    for (iop = 0; iop < nop; ++iop) {
656 1
        if (op[iop] != NULL) {
657
            PyArrayObject *ao;
658 1
            int fromanyflags = 0;
659

660 1
            if (op_flags[iop]&(NPY_ITER_READWRITE|NPY_ITER_WRITEONLY)) {
661 1
                fromanyflags |= NPY_ARRAY_WRITEBACKIFCOPY;
662
            }
663 1
            ao = (PyArrayObject *)PyArray_FROM_OF((PyObject *)op[iop],
664
                                                  fromanyflags);
665 1
            if (ao == NULL) {
666 1
                if (PyErr_Occurred() &&
667 1
                            PyErr_ExceptionMatches(PyExc_TypeError)) {
668 1
                    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 1
                for (iop = 0; iop < nop; ++iop) {
674 1
                    Py_DECREF(op[iop]);
675
                }
676 1
                *nop_out = 0;
677 1
                return 0;
678
            }
679 1
            Py_DECREF(op[iop]);
680 1
            op[iop] = ao;
681
        }
682
    }
683

684
    return 1;
685
}
686

687
static int
688 1
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 1
    PyObject *op_in = NULL, *op_flags_in = NULL,
696 1
                *op_dtypes_in = NULL, *op_axes_in = NULL;
697

698 1
    int iop, nop = 0;
699
    PyArrayObject *op[NPY_MAXARGS];
700 1
    npy_uint32 flags = 0;
701 1
    NPY_ORDER order = NPY_KEEPORDER;
702 1
    NPY_CASTING casting = NPY_SAFE_CASTING;
703
    npy_uint32 op_flags[NPY_MAXARGS];
704
    PyArray_Descr *op_request_dtypes[NPY_MAXARGS];
705 1
    int oa_ndim = -1;
706
    int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS];
707
    int *op_axes[NPY_MAXARGS];
708 1
    PyArray_Dims itershape = {NULL, -1};
709 1
    int buffersize = 0;
710

711 1
    if (self->iter != NULL) {
712 0
        PyErr_SetString(PyExc_ValueError,
713
                "Iterator was already initialized");
714 0
        return -1;
715
    }
716

717 1
    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 1
        npy_free_cache_dim_obj(itershape);
728 1
        return -1;
729
    }
730

731
    /* Set the dtypes and ops to all NULL to start */
732 1
    memset(op_request_dtypes, 0, sizeof(op_request_dtypes));
733

734
    /* op and op_flags */
735 1
    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 1
    if (op_dtypes_in != NULL && op_dtypes_in != Py_None &&
742 1
            npyiter_convert_dtypes(op_dtypes_in,
743
                                   op_request_dtypes, nop) != 1) {
744
        goto fail;
745
    }
746

747
    /* op_axes */
748 1
    if (op_axes_in != NULL && op_axes_in != Py_None) {
749
        /* Initialize to point to the op_axes arrays */
750 1
        for (iop = 0; iop < nop; ++iop) {
751 1
            op_axes[iop] = op_axes_arrays[iop];
752
        }
753

754 1
        if (npyiter_convert_op_axes(op_axes_in, nop,
755
                                    op_axes, &oa_ndim) != 1) {
756
            goto fail;
757
        }
758
    }
759

760 1
    if (itershape.len != -1) {
761 1
        if (oa_ndim == -1) {
762 1
            oa_ndim = itershape.len;
763 1
            memset(op_axes, 0, sizeof(op_axes[0]) * nop);
764
        }
765 1
        else if (oa_ndim != itershape.len) {
766 0
            PyErr_SetString(PyExc_ValueError,
767
                        "'op_axes' and 'itershape' must have the same number "
768
                        "of entries equal to the iterator ndim");
769 0
            goto fail;
770
        }
771
    }
772

773 1
    self->iter = NpyIter_AdvancedNew(nop, op, flags, order, casting, op_flags,
774
                                  op_request_dtypes,
775 1
                                  oa_ndim, oa_ndim >= 0 ? op_axes : NULL,
776
                                  itershape.ptr,
777
                                  buffersize);
778

779 1
    if (self->iter == NULL) {
780
        goto fail;
781
    }
782

783
    /* Cache some values for the member functions to use */
784 1
    if (npyiter_cache_values(self) < 0) {
785
        goto fail;
786
    }
787

788 1
    if (NpyIter_GetIterSize(self->iter) == 0) {
789 1
        self->started = 1;
790 1
        self->finished = 1;
791
    }
792
    else {
793 1
        self->started = 0;
794 1
        self->finished = 0;
795
    }
796

797 1
    npy_free_cache_dim_obj(itershape);
798

799
    /* Release the references we got to the ops and dtypes */
800 1
    for (iop = 0; iop < nop; ++iop) {
801 1
        Py_XDECREF(op[iop]);
802 1
        Py_XDECREF(op_request_dtypes[iop]);
803
    }
804

805
    return 0;
806

807 1
fail:
808 1
    npy_free_cache_dim_obj(itershape);
809 1
    for (iop = 0; iop < nop; ++iop) {
810 1
        Py_XDECREF(op[iop]);
811 1
        Py_XDECREF(op_request_dtypes[iop]);
812
    }
813
    return -1;
814
}
815

816
NPY_NO_EXPORT PyObject *
817 1
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 1
    PyObject *op_in = NULL, *axes_in = NULL,
826 1
            *op_flags_in = NULL, *op_dtypes_in = NULL;
827

828 1
    int iop, nop = 0, inest, nnest = 0;
829
    PyArrayObject *op[NPY_MAXARGS];
830 1
    npy_uint32 flags = 0, flags_inner;
831 1
    NPY_ORDER order = NPY_KEEPORDER;
832 1
    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 1
    int buffersize = 0;
842

843 1
    PyObject *ret = NULL;
844

845 1
    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 1
    if (!PyTuple_Check(axes_in) && !PyList_Check(axes_in)) {
859 0
        PyErr_SetString(PyExc_ValueError,
860
                "axes must be a tuple of axis arrays");
861 0
        return NULL;
862
    }
863 1
    nnest = PySequence_Size(axes_in);
864 1
    if (nnest < 2) {
865 0
        PyErr_SetString(PyExc_ValueError,
866
                "axes must have at least 2 entries for nested iteration");
867 0
        return NULL;
868
    }
869 1
    naxes = 0;
870 1
    memset(used_axes, 0, NPY_MAXDIMS);
871 1
    for (inest = 0; inest < nnest; ++inest) {
872 1
        PyObject *item = PySequence_GetItem(axes_in, inest);
873
        npy_intp i;
874 1
        if (item == NULL) {
875
            return NULL;
876
        }
877 1
        if (!PyTuple_Check(item) && !PyList_Check(item)) {
878 0
            PyErr_SetString(PyExc_ValueError,
879
                    "Each item in axes must be a an integer tuple");
880 0
            Py_DECREF(item);
881
            return NULL;
882
        }
883 1
        nested_naxes[inest] = PySequence_Size(item);
884 1
        if (naxes + nested_naxes[inest] > NPY_MAXDIMS) {
885 0
            PyErr_SetString(PyExc_ValueError,
886
                    "Too many axes given");
887 0
            Py_DECREF(item);
888
            return NULL;
889
        }
890 1
        for (i = 0; i < nested_naxes[inest]; ++i) {
891 1
            PyObject *v = PySequence_GetItem(item, i);
892
            npy_intp axis;
893 1
            if (v == NULL) {
894 0
                Py_DECREF(item);
895
                return NULL;
896
            }
897 1
            axis = PyLong_AsLong(v);
898 1
            Py_DECREF(v);
899 1
            if (axis < 0 || axis >= NPY_MAXDIMS) {
900 0
                PyErr_SetString(PyExc_ValueError,
901
                        "An axis is out of bounds");
902 0
                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 1
            if (used_axes[axis] != 0) {
910 0
                PyErr_SetString(PyExc_ValueError,
911
                        "An axis is used more than once");
912 0
                Py_DECREF(item);
913
                return NULL;
914
            }
915 1
            used_axes[axis] = 1;
916 1
            op_axes_data[naxes+i] = axis;
917
        }
918 1
        nested_op_axes[inest] = &op_axes_data[naxes];
919 1
        naxes += nested_naxes[inest];
920 1
        Py_DECREF(item);
921
    }
922

923
    /* op and op_flags */
924 1
    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 1
    memset(op_request_dtypes, 0, sizeof(op_request_dtypes[0])*nop);
931 1
    memset(op_request_dtypes_inner, 0,
932
                        sizeof(op_request_dtypes_inner[0])*nop);
933

934
    /* op_request_dtypes */
935 1
    if (op_dtypes_in != NULL && op_dtypes_in != Py_None &&
936 1
            npyiter_convert_dtypes(op_dtypes_in,
937
                                   op_request_dtypes, nop) != 1) {
938
        goto fail;
939
    }
940

941 1
    ret = PyTuple_New(nnest);
942 1
    if (ret == NULL) {
943
        goto fail;
944
    }
945

946
    /* For broadcasting allocated arrays */
947 1
    for (iaxes = 0; iaxes < naxes; ++iaxes) {
948 1
        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 1
    for (iop = 0; iop < nop; ++iop) {
957 1
        if ((op_flags[iop]&NPY_ITER_ALLOCATE) && op[iop] != NULL) {
958 0
            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 1
        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 1
        if ((flags&(NPY_ITER_BUFFERED)) &&
974 1
                !(op_flags[iop]&(NPY_ITER_COPY|
975
                                   NPY_ITER_UPDATEIFCOPY|
976
                                   NPY_ITER_ALLOCATE))) {
977 1
            op_flags[iop] &= ~(NPY_ITER_NBO|NPY_ITER_ALIGNED|NPY_ITER_CONTIG);
978 1
            op_request_dtypes_inner[iop] = op_request_dtypes[iop];
979 1
            op_request_dtypes[iop] = NULL;
980
        }
981
    }
982

983
    /* Only the inner loop gets the buffering and no inner flags */
984 1
    flags_inner = flags&~NPY_ITER_COMMON_DTYPE;
985 1
    flags &= ~(NPY_ITER_EXTERNAL_LOOP|
986
                    NPY_ITER_BUFFERED);
987

988 1
    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 1
        for (iop = 0; iop < nop; ++iop) {
997 1
            if (op_flags[iop]&NPY_ITER_ALLOCATE) {
998 0
                if (inest == 0) {
999 0
                    op_axes_nop[iop] = NULL;
1000
                }
1001
                else {
1002 0
                    op_axes_nop[iop] = negones;
1003
                }
1004
            }
1005
            else {
1006 1
                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 1
        iter = (NewNpyArrayIterObject *)npyiter_new(&NpyIter_Type, NULL, NULL);
1024 1
        if (iter == NULL) {
1025 0
            Py_DECREF(ret);
1026 0
            goto fail;
1027
        }
1028

1029 1
        if (inest < nnest-1) {
1030 1
            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 1
            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 1
        if (iter->iter == NULL) {
1046 0
            Py_DECREF(ret);
1047
            goto fail;
1048
        }
1049

1050
        /* Cache some values for the member functions to use */
1051 1
        if (npyiter_cache_values(iter) < 0) {
1052 0
            Py_DECREF(ret);
1053
            goto fail;
1054
        }
1055

1056 1
        if (NpyIter_GetIterSize(iter->iter) == 0) {
1057 0
            iter->started = 1;
1058 0
            iter->finished = 1;
1059
        }
1060
        else {
1061 1
            iter->started = 0;
1062 1
            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 1
        if (inest == 0) {
1070 1
            PyArrayObject **operands = NpyIter_GetOperandArray(iter->iter);
1071 1
            for (iop = 0; iop < nop; ++iop) {
1072 1
                if (op[iop] != operands[iop]) {
1073 1
                    Py_XDECREF(op[iop]);
1074 1
                    op[iop] = operands[iop];
1075 1
                    Py_INCREF(op[iop]);
1076
                }
1077

1078
                /*
1079
                 * Clear any flags allowing copies for
1080
                 * the rest of the iterators
1081
                 */
1082 1
                op_flags[iop] &= ~(NPY_ITER_COPY|
1083
                                 NPY_ITER_UPDATEIFCOPY);
1084
            }
1085
            /* Clear the common dtype flag for the rest of the iterators */
1086 1
            flags &= ~NPY_ITER_COMMON_DTYPE;
1087
        }
1088

1089 1
        PyTuple_SET_ITEM(ret, inest, (PyObject *)iter);
1090
    }
1091

1092
    /* Release our references to the ops and dtypes */
1093 1
    for (iop = 0; iop < nop; ++iop) {
1094 1
        Py_XDECREF(op[iop]);
1095 1
        Py_XDECREF(op_request_dtypes[iop]);
1096 1
        Py_XDECREF(op_request_dtypes_inner[iop]);
1097
    }
1098

1099
    /* Set up the nested child references */
1100 1
    for (inest = 0; inest < nnest-1; ++inest) {
1101
        NewNpyArrayIterObject *iter;
1102 1
        iter = (NewNpyArrayIterObject *)PyTuple_GET_ITEM(ret, inest);
1103
        /*
1104
         * Indicates which iterator to reset with new base pointers
1105
         * each iteration step.
1106
         */
1107 1
        iter->nested_child =
1108 1
                (NewNpyArrayIterObject *)PyTuple_GET_ITEM(ret, inest+1);
1109 1
        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 1
        if (NpyIter_ResetBasePointers(iter->nested_child->iter,
1115
                                iter->dataptrs, NULL) != NPY_SUCCEED) {
1116 0
            Py_DECREF(ret);
1117
            return NULL;
1118
        }
1119
    }
1120

1121
    return ret;
1122

1123 0
fail:
1124 0
    for (iop = 0; iop < nop; ++iop) {
1125 0
        Py_XDECREF(op[iop]);
1126 0
        Py_XDECREF(op_request_dtypes[iop]);
1127 0
        Py_XDECREF(op_request_dtypes_inner[iop]);
1128
    }
1129
    return NULL;
1130
}
1131

1132

1133
static void
1134 1
npyiter_dealloc(NewNpyArrayIterObject *self)
1135
{
1136 1
    if (self->iter) {
1137 1
        if (npyiter_has_writeback(self->iter)) {
1138 1
            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 0
                s = PyUnicode_FromString("npyiter_dealloc");
1146 0
                if (s) {
1147 0
                    PyErr_WriteUnraisable(s);
1148 0
                    Py_DECREF(s);
1149
                }
1150
                else {
1151 0
                    PyErr_WriteUnraisable(Py_None);
1152
                }
1153
            }
1154
        }
1155 1
        NpyIter_Deallocate(self->iter);
1156 1
        self->iter = NULL;
1157 1
        Py_XDECREF(self->nested_child);
1158 1
        self->nested_child = NULL;
1159
    }
1160 1
    Py_TYPE(self)->tp_free((PyObject*)self);
1161
}
1162

1163
static int
1164 1
npyiter_resetbasepointers(NewNpyArrayIterObject *self)
1165
{
1166 1
    while (self->nested_child) {
1167 1
        if (NpyIter_ResetBasePointers(self->nested_child->iter,
1168
                                        self->dataptrs, NULL) != NPY_SUCCEED) {
1169
            return NPY_FAIL;
1170
        }
1171 1
        self = self->nested_child;
1172 1
        if (NpyIter_GetIterSize(self->iter) == 0) {
1173 0
            self->started = 1;
1174 0
            self->finished = 1;
1175
        }
1176
        else {
1177 1
            self->started = 0;
1178 1
            self->finished = 0;
1179
        }
1180
    }
1181

1182
    return NPY_SUCCEED;
1183
}
1184

1185
static PyObject *
1186 1
npyiter_reset(NewNpyArrayIterObject *self)
1187
{
1188 1
    if (self->iter == NULL) {
1189 0
        PyErr_SetString(PyExc_ValueError,
1190
                "Iterator is invalid");
1191 0
        return NULL;
1192
    }
1193

1194 1
    if (NpyIter_Reset(self->iter, NULL) != NPY_SUCCEED) {
1195
        return NULL;
1196
    }
1197 1
    if (NpyIter_GetIterSize(self->iter) == 0) {
1198 0
        self->started = 1;
1199 0
        self->finished = 1;
1200
    }
1201
    else {
1202 1
        self->started = 0;
1203 1
        self->finished = 0;
1204
    }
1205

1206 1
    if (self->get_multi_index == NULL && NpyIter_HasMultiIndex(self->iter)) {
1207 1
        self->get_multi_index = NpyIter_GetGetMultiIndex(self->iter, NULL);
1208
    }
1209

1210
    /* If there is nesting, the nested iterators should be reset */
1211 1
    if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1212
        return NULL;
1213
    }
1214

1215 1
    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 *
1223 1
npyiter_copy(NewNpyArrayIterObject *self)
1224
{
1225
    NewNpyArrayIterObject *iter;
1226

1227 1
    if (self->iter == NULL) {
1228 0
        PyErr_SetString(PyExc_ValueError,
1229
                "Iterator is invalid");
1230 0
        return NULL;
1231
    }
1232

1233
    /* Allocate the iterator */
1234 1
    iter = (NewNpyArrayIterObject *)npyiter_new(&NpyIter_Type, NULL, NULL);
1235 1
    if (iter == NULL) {
1236
        return NULL;
1237
    }
1238

1239
    /* Copy the C iterator */
1240 1
    iter->iter = NpyIter_Copy(self->iter);
1241 1
    if (iter->iter == NULL) {
1242 0
        Py_DECREF(iter);
1243
        return NULL;
1244
    }
1245

1246
    /* Cache some values for the member functions to use */
1247 1
    if (npyiter_cache_values(iter) < 0) {
1248 0
        Py_DECREF(iter);
1249
        return NULL;
1250
    }
1251

1252 1
    iter->started = self->started;
1253 1
    iter->finished = self->finished;
1254

1255 1
    return (PyObject *)iter;
1256
}
1257

1258
static PyObject *
1259 1
npyiter_iternext(NewNpyArrayIterObject *self)
1260
{
1261 1
    if (self->iter != NULL && self->iternext != NULL &&
1262 1
                        !self->finished && self->iternext(self->iter)) {
1263
        /* If there is nesting, the nested iterators should be reset */
1264 1
        if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1265
            return NULL;
1266
        }
1267

1268 1
        Py_RETURN_TRUE;
1269
    }
1270
    else {
1271 1
        if (PyErr_Occurred()) {
1272
            /* casting error, buffer cleanup will occur at reset or dealloc */
1273
            return NULL;
1274
        }
1275 1
        self->finished = 1;
1276 1
        Py_RETURN_FALSE;
1277
    }
1278
}
1279

1280
static PyObject *
1281 1
npyiter_remove_axis(NewNpyArrayIterObject *self, PyObject *args)
1282
{
1283 1
    int axis = 0;
1284

1285 1
    if (self->iter == NULL) {
1286 0
        PyErr_SetString(PyExc_ValueError,
1287
                "Iterator is invalid");
1288 0
        return NULL;
1289
    }
1290

1291 1
    if (!PyArg_ParseTuple(args, "i:remove_axis", &axis)) {
1292
        return NULL;
1293
    }
1294

1295 1
    if (NpyIter_RemoveAxis(self->iter, axis) != NPY_SUCCEED) {
1296
        return NULL;
1297
    }
1298
    /* RemoveAxis invalidates cached values */
1299 1
    if (npyiter_cache_values(self) < 0) {
1300
        return NULL;
1301
    }
1302
    /* RemoveAxis also resets the iterator */
1303 1
    if (NpyIter_GetIterSize(self->iter) == 0) {
1304 0
        self->started = 1;
1305 0
        self->finished = 1;
1306
    }
1307
    else {
1308 1
        self->started = 0;
1309 1
        self->finished = 0;
1310
    }
1311

1312 1
    Py_RETURN_NONE;
1313
}
1314

1315
static PyObject *
1316 1
npyiter_remove_multi_index(NewNpyArrayIterObject *self)
1317
{
1318 1
    if (self->iter == NULL) {
1319 0
        PyErr_SetString(PyExc_ValueError,
1320
                "Iterator is invalid");
1321 0
        return NULL;
1322
    }
1323

1324 1
    NpyIter_RemoveMultiIndex(self->iter);
1325
    /* RemoveMultiIndex invalidates cached values */
1326 1
    npyiter_cache_values(self);
1327
    /* RemoveMultiIndex also resets the iterator */
1328 1
    if (NpyIter_GetIterSize(self->iter) == 0) {
1329 0
        self->started = 1;
1330 0
        self->finished = 1;
1331
    }
1332
    else {
1333 1
        self->started = 0;
1334 1
        self->finished = 0;
1335
    }
1336

1337 1
    Py_RETURN_NONE;
1338
}
1339

1340
static PyObject *
1341 1
npyiter_enable_external_loop(NewNpyArrayIterObject *self)
1342
{
1343 1
    if (self->iter == NULL) {
1344 0
        PyErr_SetString(PyExc_ValueError,
1345
                "Iterator is invalid");
1346 0
        return NULL;
1347
    }
1348

1349 1
    NpyIter_EnableExternalLoop(self->iter);
1350
    /* EnableExternalLoop invalidates cached values */
1351 1
    npyiter_cache_values(self);
1352
    /* EnableExternalLoop also resets the iterator */
1353 1
    if (NpyIter_GetIterSize(self->iter) == 0) {
1354 0
        self->started = 1;
1355 0
        self->finished = 1;
1356
    }
1357
    else {
1358 1
        self->started = 0;
1359 1
        self->finished = 0;
1360
    }
1361

1362 1
    Py_RETURN_NONE;
1363
}
1364

1365
static PyObject *
1366 0
npyiter_debug_print(NewNpyArrayIterObject *self)
1367
{
1368 0
    if (self->iter != NULL) {
1369 0
        NpyIter_DebugPrint(self->iter);
1370
    }
1371
    else {
1372
        printf("Iterator: (nil)\n");
1373
    }
1374

1375 0
    Py_RETURN_NONE;
1376
}
1377

1378
NPY_NO_EXPORT PyObject *
1379
npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i);
1380

1381 1
static PyObject *npyiter_value_get(NewNpyArrayIterObject *self)
1382
{
1383
    PyObject *ret;
1384

1385
    npy_intp iop, nop;
1386

1387 1
    if (self->iter == NULL || self->finished) {
1388 0
        PyErr_SetString(PyExc_ValueError,
1389
                "Iterator is past the end");
1390 0
        return NULL;
1391
    }
1392

1393 1
    nop = NpyIter_GetNOp(self->iter);
1394

1395
    /* Return an array  or tuple of arrays with the values */
1396 1
    if (nop == 1) {
1397 1
        ret = npyiter_seq_item(self, 0);
1398
    }
1399
    else {
1400 1
        ret = PyTuple_New(nop);
1401 1
        if (ret == NULL) {
1402
            return NULL;
1403
        }
1404 1
        for (iop = 0; iop < nop; ++iop) {
1405 1
            PyObject *a = npyiter_seq_item(self, iop);
1406 1
            if (a == NULL) {
1407 0
                Py_DECREF(ret);
1408
                return NULL;
1409
            }
1410 1
            PyTuple_SET_ITEM(ret, iop, a);
1411
        }
1412
    }
1413

1414
    return ret;
1415
}
1416

1417 1
static PyObject *npyiter_operands_get(NewNpyArrayIterObject *self)
1418
{
1419
    PyObject *ret;
1420

1421
    npy_intp iop, nop;
1422
    PyArrayObject **operands;
1423

1424 1
    if (self->iter == NULL) {
1425 1
        PyErr_SetString(PyExc_ValueError,
1426
                "Iterator is invalid");
1427 1
        return NULL;
1428
    }
1429 1
    nop = NpyIter_GetNOp(self->iter);
1430 1
    operands = self->operands;
1431

1432 1
    ret = PyTuple_New(nop);
1433 1
    if (ret == NULL) {
1434
        return NULL;
1435
    }
1436 1
    for (iop = 0; iop < nop; ++iop) {
1437 1
        PyObject *operand = (PyObject *)operands[iop];
1438

1439 1
        Py_INCREF(operand);
1440 1
        PyTuple_SET_ITEM(ret, iop, operand);
1441
    }
1442

1443
    return ret;
1444
}
1445

1446 1
static PyObject *npyiter_itviews_get(NewNpyArrayIterObject *self)
1447
{
1448
    PyObject *ret;
1449

1450
    npy_intp iop, nop;
1451

1452 1
    if (self->iter == NULL) {
1453 0
        PyErr_SetString(PyExc_ValueError,
1454
                "Iterator is invalid");
1455 0
        return NULL;
1456
    }
1457 1
    nop = NpyIter_GetNOp(self->iter);
1458

1459 1
    ret = PyTuple_New(nop);
1460 1
    if (ret == NULL) {
1461
        return NULL;
1462
    }
1463 1
    for (iop = 0; iop < nop; ++iop) {
1464 1
        PyArrayObject *view = NpyIter_GetIterView(self->iter, iop);
1465

1466 1
        if (view == NULL) {
1467 0
            Py_DECREF(ret);
1468
            return NULL;
1469
        }
1470 1
        PyTuple_SET_ITEM(ret, iop, (PyObject *)view);
1471
    }
1472

1473
    return ret;
1474
}
1475

1476
static PyObject *
1477 1
npyiter_next(NewNpyArrayIterObject *self)
1478
{
1479 1
    if (self->iter == NULL || self->iternext == NULL ||
1480 1
                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 1
    if (self->started) {
1489 1
        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 1
            self->finished = 1;
1495 1
            return NULL;
1496
        }
1497

1498
        /* If there is nesting, the nested iterators should be reset */
1499 1
        if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1500
            return NULL;
1501
        }
1502
    }
1503 1
    self->started = 1;
1504

1505 1
    return npyiter_value_get(self);
1506
};
1507

1508 1
static PyObject *npyiter_shape_get(NewNpyArrayIterObject *self)
1509
{
1510
    PyObject *ret;
1511
    npy_intp idim, ndim, shape[NPY_MAXDIMS];
1512

1513 1
    if (self->iter == NULL || self->finished) {
1514 1
        PyErr_SetString(PyExc_ValueError,
1515
                "Iterator is past the end");
1516 1
        return NULL;
1517
    }
1518

1519 1
    if (NpyIter_GetShape(self->iter, shape) == NPY_SUCCEED) {
1520 1
        ndim = NpyIter_GetNDim(self->iter);
1521 1
        ret = PyTuple_New(ndim);
1522 1
        if (ret != NULL) {
1523 1
            for (idim = 0; idim < ndim; ++idim) {
1524 1
                PyTuple_SET_ITEM(ret, idim,
1525
                        PyLong_FromLong(shape[idim]));
1526
            }
1527
            return ret;
1528
        }
1529
    }
1530

1531
    return NULL;
1532
}
1533

1534 1
static PyObject *npyiter_multi_index_get(NewNpyArrayIterObject *self)
1535
{
1536
    PyObject *ret;
1537
    npy_intp idim, ndim, multi_index[NPY_MAXDIMS];
1538

1539 1
    if (self->iter == NULL || self->finished) {
1540 0
        PyErr_SetString(PyExc_ValueError,
1541
                "Iterator is past the end");
1542 0
        return NULL;
1543
    }
1544

1545 1
    if (self->get_multi_index != NULL) {
1546 1
        ndim = NpyIter_GetNDim(self->iter);
1547 1
        self->get_multi_index(self->iter, multi_index);
1548 1
        ret = PyTuple_New(ndim);
1549 1
        if (ret == NULL) {
1550
            return NULL;
1551
        }
1552 1
        for (idim = 0; idim < ndim; ++idim) {
1553 1
            PyTuple_SET_ITEM(ret, idim,
1554
                    PyLong_FromLong(multi_index[idim]));
1555
        }
1556
        return ret;
1557
    }
1558
    else {
1559 1
        if (!NpyIter_HasMultiIndex(self->iter)) {
1560 1
            PyErr_SetString(PyExc_ValueError,
1561
                    "Iterator is not tracking a multi-index");
1562 1
            return NULL;
1563
        }
1564 1
        else if (NpyIter_HasDelayedBufAlloc(self->iter)) {
1565 1
            PyErr_SetString(PyExc_ValueError,
1566
                    "Iterator construction used delayed buffer allocation, "
1567
                    "and no reset has been done yet");
1568 1
            return NULL;
1569
        }
1570
        else {
1571 0
            PyErr_SetString(PyExc_ValueError,
1572
                    "Iterator is in an invalid state");
1573 0
            return NULL;
1574
        }
1575
    }
1576
}
1577

1578
static int
1579 1
npyiter_multi_index_set(NewNpyArrayIterObject *self, PyObject *value)
1580
{
1581
    npy_intp idim, ndim, multi_index[NPY_MAXDIMS];
1582

1583 1
    if (value == NULL) {
1584 1
        PyErr_SetString(PyExc_AttributeError,
1585
                "Cannot delete nditer multi_index");
1586 1
        return -1;
1587
    }
1588 1
    if (self->iter == NULL) {
1589 0
        PyErr_SetString(PyExc_ValueError,
1590
                "Iterator is invalid");
1591 0
        return -1;
1592
    }
1593

1594 1
    if (NpyIter_HasMultiIndex(self->iter)) {
1595 0
        ndim = NpyIter_GetNDim(self->iter);
1596 0
        if (!PySequence_Check(value)) {
1597 0
            PyErr_SetString(PyExc_ValueError,
1598
                    "multi_index must be set with a sequence");
1599 0
            return -1;
1600
        }
1601 0
        if (PySequence_Size(value) != ndim) {
1602 0
            PyErr_SetString(PyExc_ValueError,
1603
                    "Wrong number of indices");
1604 0
            return -1;
1605
        }
1606 0
        for (idim = 0; idim < ndim; ++idim) {
1607 0
            PyObject *v = PySequence_GetItem(value, idim);
1608 0
            multi_index[idim] = PyLong_AsLong(v);
1609 0
            if (error_converting(multi_index[idim])) {
1610 0
                Py_XDECREF(v);
1611
                return -1;
1612
            }
1613
        }
1614 0
        if (NpyIter_GotoMultiIndex(self->iter, multi_index) != NPY_SUCCEED) {
1615
            return -1;
1616
        }
1617 0
        self->started = 0;
1618 0
        self->finished = 0;
1619

1620
        /* If there is nesting, the nested iterators should be reset */
1621 0
        if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1622
            return -1;
1623
        }
1624

1625 0
        return 0;
1626
    }
1627
    else {
1628 1
        PyErr_SetString(PyExc_ValueError,
1629
                "Iterator is not tracking a multi-index");
1630 1
        return -1;
1631
    }
1632
}
1633

1634 1
static PyObject *npyiter_index_get(NewNpyArrayIterObject *self)
1635
{
1636 1
    if (self->iter == NULL || self->finished) {
1637 0
        PyErr_SetString(PyExc_ValueError,
1638
                "Iterator is past the end");
1639 0
        return NULL;
1640
    }
1641

1642 1
    if (NpyIter_HasIndex(self->iter)) {
1643 1
        npy_intp ind = *NpyIter_GetIndexPtr(self->iter);
1644 1
        return PyLong_FromLong(ind);
1645
    }
1646
    else {
1647 1
        PyErr_SetString(PyExc_ValueError,
1648
                "Iterator does not have an index");
1649 1
        return NULL;
1650
    }
1651
}
1652

1653 1
static int npyiter_index_set(NewNpyArrayIterObject *self, PyObject *value)
1654
{
1655 1
    if (value == NULL) {
1656 1
        PyErr_SetString(PyExc_AttributeError,
1657
                "Cannot delete nditer index");
1658 1
        return -1;
1659
    }
1660 1
    if (self->iter == NULL) {
1661 0
        PyErr_SetString(PyExc_ValueError,
1662
                "Iterator is invalid");
1663 0
        return -1;
1664
    }
1665

1666 1
    if (NpyIter_HasIndex(self->iter)) {
1667
        npy_intp ind;
1668 0
        ind = PyLong_AsLong(value);
1669 0
        if (error_converting(ind)) {
1670
            return -1;
1671
        }
1672 0
        if (NpyIter_GotoIndex(self->iter, ind) != NPY_SUCCEED) {
1673
            return -1;
1674
        }
1675 0
        self->started = 0;
1676 0
        self->finished = 0;
1677

1678
        /* If there is nesting, the nested iterators should be reset */
1679 0
        if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1680
            return -1;
1681
        }
1682

1683 0
        return 0;
1684
    }
1685
    else {
1686 1
        PyErr_SetString(PyExc_ValueError,
1687
                "Iterator does not have an index");
1688 1
        return -1;
1689
    }
1690
}
1691

1692 1
static PyObject *npyiter_iterindex_get(NewNpyArrayIterObject *self)
1693
{
1694 1
    if (self->iter == NULL || self->finished) {
1695 0
        PyErr_SetString(PyExc_ValueError,
1696
                "Iterator is past the end");
1697 0
        return NULL;
1698
    }
1699

1700 1
    return PyLong_FromLong(NpyIter_GetIterIndex(self->iter));
1701
}
1702

1703 1
static int npyiter_iterindex_set(NewNpyArrayIterObject *self, PyObject *value)
1704
{
1705
    npy_intp iterindex;
1706

1707 1
    if (value == NULL) {
1708 1
        PyErr_SetString(PyExc_AttributeError,
1709
                "Cannot delete nditer iterindex");
1710 1
        return -1;
1711
    }
1712 1
    if (self->iter == NULL) {
1713 0
        PyErr_SetString(PyExc_ValueError,
1714
                "Iterator is invalid");
1715 0
        return -1;
1716
    }
1717

1718 1
    iterindex = PyLong_AsLong(value);
1719 1
    if (error_converting(iterindex)) {
1720
        return -1;
1721
    }
1722 1
    if (NpyIter_GotoIterIndex(self->iter, iterindex) != NPY_SUCCEED) {
1723
        return -1;
1724
    }
1725 1
    self->started = 0;
1726 1
    self->finished = 0;
1727

1728
    /* If there is nesting, the nested iterators should be reset */
1729 1
    if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1730
        return -1;
1731
    }
1732

1733 1
    return 0;
1734
}
1735

1736 1
static PyObject *npyiter_iterrange_get(NewNpyArrayIterObject *self)
1737
{
1738 1
    npy_intp istart = 0, iend = 0;
1739
    PyObject *ret;
1740

1741 1
    if (self->iter == NULL) {
1742 0
        PyErr_SetString(PyExc_ValueError,
1743
                "Iterator is invalid");
1744 0
        return NULL;
1745
    }
1746

1747 1
    NpyIter_GetIterIndexRange(self->iter, &istart, &iend);
1748

1749 1
    ret = PyTuple_New(2);
1750 1
    if (ret == NULL) {
1751
        return NULL;
1752
    }
1753

1754 1
    PyTuple_SET_ITEM(ret, 0, PyLong_FromLong(istart));
1755 1
    PyTuple_SET_ITEM(ret, 1, PyLong_FromLong(iend));
1756

1757 1
    return ret;
1758
}
1759

1760 1
static int npyiter_iterrange_set(NewNpyArrayIterObject *self, PyObject *value)
1761
{
1762 1
    npy_intp istart = 0, iend = 0;
1763

1764 1
    if (value == NULL) {
1765 1
        PyErr_SetString(PyExc_AttributeError,
1766
                "Cannot delete nditer iterrange");
1767 1
        return -1;
1768
    }
1769 1
    if (self->iter == NULL) {
1770 0
        PyErr_SetString(PyExc_ValueError,
1771
                "Iterator is invalid");
1772 0
        return -1;
1773
    }
1774

1775 1
    if (!PyArg_ParseTuple(value, "nn", &istart, &iend)) {
1776
        return -1;
1777
    }
1778

1779 1
    if (NpyIter_ResetToIterIndexRange(self->iter, istart, iend, NULL)
1780
                                                    != NPY_SUCCEED) {
1781
        return -1;
1782
    }
1783 1
    if (istart < iend) {
1784 1
        self->started = self->finished = 0;
1785
    }
1786
    else {
1787 1
        self->started = self->finished = 1;
1788
    }
1789

1790 1
    if (self->get_multi_index == NULL && NpyIter_HasMultiIndex(self->iter)) {
1791 0
        self->get_multi_index = NpyIter_GetGetMultiIndex(self->iter, NULL);
1792
    }
1793

1794
    /* If there is nesting, the nested iterators should be reset */
1795 1
    if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
1796
        return -1;
1797
    }
1798

1799 1
    return 0;
1800
}
1801

1802 1
static PyObject *npyiter_has_delayed_bufalloc_get(NewNpyArrayIterObject *self)
1803
{
1804 1
    if (self->iter == NULL) {
1805 0
        PyErr_SetString(PyExc_ValueError,
1806
                "Iterator is invalid");
1807 0
        return NULL;
1808
    }
1809

1810 1
    if (NpyIter_HasDelayedBufAlloc(self->iter)) {
1811 1
        Py_RETURN_TRUE;
1812
    }
1813
    else {
1814 1
        Py_RETURN_FALSE;
1815
    }
1816
}
1817

1818 1
static PyObject *npyiter_iterationneedsapi_get(NewNpyArrayIterObject *self)
1819
{
1820 1
    if (self->iter == NULL) {
1821 0
        PyErr_SetString(PyExc_ValueError,
1822
                "Iterator is invalid");
1823 0
        return NULL;
1824
    }
1825

1826 1
    if (NpyIter_IterationNeedsAPI(self->iter)) {
1827 1
        Py_RETURN_TRUE;
1828
    }
1829
    else {
1830 1
        Py_RETURN_FALSE;
1831
    }
1832
}
1833

1834 0
static PyObject *npyiter_has_multi_index_get(NewNpyArrayIterObject *self)
1835
{
1836 0
    if (self->iter == NULL) {
1837 0
        PyErr_SetString(PyExc_ValueError,
1838
                "Iterator is invalid");
1839 0
        return NULL;
1840
    }
1841

1842 0
    if (NpyIter_HasMultiIndex(self->iter)) {
1843 0
        Py_RETURN_TRUE;
1844
    }
1845
    else {
1846 0
        Py_RETURN_FALSE;
1847
    }
1848
}
1849

1850 0
static PyObject *npyiter_has_index_get(NewNpyArrayIterObject *self)
1851
{
1852 0
    if (self->iter == NULL) {
1853 0
        PyErr_SetString(PyExc_ValueError,
1854
                "Iterator is invalid");
1855 0
        return NULL;
1856
    }
1857

1858 0
    if (NpyIter_HasIndex(self->iter)) {
1859 0
        Py_RETURN_TRUE;
1860
    }
1861
    else {
1862 0
        Py_RETURN_FALSE;
1863
    }
1864
}
1865

1866 1
static PyObject *npyiter_dtypes_get(NewNpyArrayIterObject *self)
1867
{
1868
    PyObject *ret;
1869

1870
    npy_intp iop, nop;
1871
    PyArray_Descr **dtypes;
1872

1873 1
    if (self->iter == NULL) {
1874 0
        PyErr_SetString(PyExc_ValueError,
1875
                "Iterator is invalid");
1876 0
        return NULL;
1877
    }
1878 1
    nop = NpyIter_GetNOp(self->iter);
1879

1880 1
    ret = PyTuple_New(nop);
1881 1
    if (ret == NULL) {
1882
        return NULL;
1883
    }
1884 1
    dtypes = self->dtypes;
1885 1
    for (iop = 0; iop < nop; ++iop) {
1886 1
        PyArray_Descr *dtype = dtypes[iop];
1887

1888 1
        Py_INCREF(dtype);
1889 1
        PyTuple_SET_ITEM(ret, iop, (PyObject *)dtype);
1890
    }
1891

1892
    return ret;
1893
}
1894

1895 1
static PyObject *npyiter_ndim_get(NewNpyArrayIterObject *self)
1896
{
1897 1
    if (self->iter == NULL) {
1898 0
        PyErr_SetString(PyExc_ValueError,
1899
                "Iterator is invalid");
1900 0
        return NULL;
1901
    }
1902

1903 1
    return PyLong_FromLong(NpyIter_GetNDim(self->iter));
1904
}
1905

1906 0
static PyObject *npyiter_nop_get(NewNpyArrayIterObject *self)
1907
{
1908 0
    if (self->iter == NULL) {
1909 0
        PyErr_SetString(PyExc_ValueError,
1910
                "Iterator is invalid");
1911 0
        return NULL;
1912
    }
1913

1914 0
    return PyLong_FromLong(NpyIter_GetNOp(self->iter));
1915
}
1916

1917 1
static PyObject *npyiter_itersize_get(NewNpyArrayIterObject *self)
1918
{
1919 1
    if (self->iter == NULL) {
1920 0
        PyErr_SetString(PyExc_ValueError,
1921
                "Iterator is invalid");
1922 0
        return NULL;
1923
    }
1924

1925 1
    return PyLong_FromLong(NpyIter_GetIterSize(self->iter));
1926
}
1927

1928 1
static PyObject *npyiter_finished_get(NewNpyArrayIterObject *self)
1929
{
1930 1
    if (self->iter == NULL || !self->finished) {
1931 1
        Py_RETURN_FALSE;
1932
    }
1933
    else {
1934 1
        Py_RETURN_TRUE;
1935
    }
1936
}
1937

1938
NPY_NO_EXPORT Py_ssize_t
1939 1
npyiter_seq_length(NewNpyArrayIterObject *self)
1940
{
1941 1
    if (self->iter == NULL) {
1942
        return 0;
1943
    }
1944
    else {
1945 1
        return NpyIter_GetNOp(self->iter);
1946
    }
1947
}
1948

1949
NPY_NO_EXPORT PyObject *
1950 1
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 1
    Py_ssize_t i_orig = i;
1958

1959 1
    if (self->iter == NULL || self->finished) {
1960 0
        PyErr_SetString(PyExc_ValueError,
1961
                "Iterator is past the end");
1962 0
        return NULL;
1963
    }
1964

1965 1
    if (NpyIter_HasDelayedBufAlloc(self->iter)) {
1966 0
        PyErr_SetString(PyExc_ValueError,
1967
                "Iterator construction used delayed buffer allocation, "
1968
                "and no reset has been done yet");
1969 0
        return NULL;
1970
    }
1971 1
    nop = NpyIter_GetNOp(self->iter);
1972

1973
    /* Negative indexing */
1974 1
    if (i < 0) {
1975 0
        i += nop;
1976
    }
1977

1978 1
    if (i < 0 || i >= nop) {
1979 0
        PyErr_Format(PyExc_IndexError,
1980
                "Iterator operand index %zd is out of bounds", i_orig);
1981 0
        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 1
    dataptr = self->dataptrs[i];
2000 1
    dtype = self->dtypes[i];
2001 1
    has_external_loop = NpyIter_HasExternalLoop(self->iter);
2002

2003 1
    if (has_external_loop) {
2004 1
        innerloopsize = *self->innerloopsizeptr;
2005 1
        innerstride = self->innerstrides[i];
2006 1
        ret_ndim = 1;
2007
    }
2008
    else {
2009 1
        innerloopsize = 1;
2010 1
        innerstride = 0;
2011
        /* If the iterator is going over every element, return array scalars */
2012 1
        ret_ndim = 0;
2013
    }
2014

2015 1
    Py_INCREF(dtype);
2016 1
    return PyArray_NewFromDescrAndBase(
2017
            &PyArray_Type, dtype,
2018
            ret_ndim, &innerloopsize, &innerstride, dataptr,
2019 1
            self->writeflags[i] ? NPY_ARRAY_WRITEABLE : 0,
2020
            NULL, (PyObject *)self);
2021
}
2022

2023
NPY_NO_EXPORT PyObject *
2024 1
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 1
    if (self->iter == NULL || self->finished) {
2032 0
        PyErr_SetString(PyExc_ValueError,
2033
                "Iterator is past the end");
2034 0
        return NULL;
2035
    }
2036

2037 1
    if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2038 0
        PyErr_SetString(PyExc_ValueError,
2039
                "Iterator construction used delayed buffer allocation, "
2040
                "and no reset has been done yet");
2041 0
        return NULL;
2042
    }
2043 1
    nop = NpyIter_GetNOp(self->iter);
2044 1
    if (ilow < 0) {
2045
        ilow = 0;
2046
    }
2047 1
    else if (ilow >= nop) {
2048 0
        ilow = nop-1;
2049
    }
2050 1
    if (ihigh < ilow) {
2051
        ihigh = ilow;
2052
    }
2053 1
    else if (ihigh > nop) {
2054 0
        ihigh = nop;
2055
    }
2056

2057 1
    ret = PyTuple_New(ihigh-ilow);
2058 1
    if (ret == NULL) {
2059
        return NULL;
2060
    }
2061 1
    for (i = ilow; i < ihigh ; ++i) {
2062 1
        PyObject *item = npyiter_seq_item(self, i);
2063 1
        if (item == NULL) {
2064 0
            Py_DECREF(ret);
2065
            return NULL;
2066
        }
2067 1
        PyTuple_SET_ITEM(ret, i-ilow, item);
2068
    }
2069
    return ret;
2070
}
2071

2072
NPY_NO_EXPORT int
2073 1
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 1
    Py_ssize_t i_orig = i;
2082

2083

2084 1
    if (v == NULL) {
2085 0
        PyErr_SetString(PyExc_TypeError,
2086
                "Cannot delete iterator elements");
2087 0
        return -1;
2088
    }
2089

2090 1
    if (self->iter == NULL || self->finished) {
2091 0
        PyErr_SetString(PyExc_ValueError,
2092
                "Iterator is past the end");
2093 0
        return -1;
2094
    }
2095

2096 1
    if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2097 0
        PyErr_SetString(PyExc_ValueError,
2098
                "Iterator construction used delayed buffer allocation, "
2099
                "and no reset has been done yet");
2100 0
        return -1;
2101
    }
2102 1
    nop = NpyIter_GetNOp(self->iter);
2103

2104
    /* Negative indexing */
2105 1
    if (i < 0) {
2106 0
        i += nop;
2107
    }
2108

2109 1
    if (i < 0 || i >= nop) {
2110 0
        PyErr_Format(PyExc_IndexError,
2111
                "Iterator operand index %zd is out of bounds", i_orig);
2112 0
        return -1;
2113
    }
2114 1
    if (!self->writeflags[i]) {
2115 0
        PyErr_Format(PyExc_RuntimeError,
2116
                "Iterator operand %zd is not writeable", i_orig);
2117 0
        return -1;
2118
    }
2119

2120 1
    dataptr = self->dataptrs[i];
2121 1
    dtype = self->dtypes[i];
2122 1
    has_external_loop = NpyIter_HasExternalLoop(self->iter);
2123

2124 1
    if (has_external_loop) {
2125 0
        innerloopsize = *self->innerloopsizeptr;
2126 0
        innerstride = self->innerstrides[i];
2127
    }
2128
    else {
2129 1
        innerloopsize = 1;
2130 1
        innerstride = 0;
2131
    }
2132

2133
    /* TODO - there should be a better way than this... */
2134 1
    Py_INCREF(dtype);
2135 1
    tmp = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype,
2136
                                1, &innerloopsize,
2137
                                &innerstride, dataptr,
2138
                                NPY_ARRAY_WRITEABLE, NULL);
2139 1
    if (tmp == NULL) {
2140
        return -1;
2141
    }
2142

2143 1
    ret = PyArray_CopyObject(tmp, v);
2144 1
    Py_DECREF(tmp);
2145
    return ret;
2146
}
2147

2148
static int
2149 1
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 1
    if (v == NULL) {
2156 0
        PyErr_SetString(PyExc_TypeError,
2157
                "Cannot delete iterator elements");
2158 0
        return -1;
2159
    }
2160

2161 1
    if (self->iter == NULL || self->finished) {
2162 0
        PyErr_SetString(PyExc_ValueError,
2163
                "Iterator is past the end");
2164 0
        return -1;
2165
    }
2166

2167 1
    if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2168 0
        PyErr_SetString(PyExc_ValueError,
2169
                "Iterator construction used delayed buffer allocation, "
2170
                "and no reset has been done yet");
2171 0
        return -1;
2172
    }
2173 1
    nop = NpyIter_GetNOp(self->iter);
2174 1
    if (ilow < 0) {
2175
        ilow = 0;
2176
    }
2177 1
    else if (ilow >= nop) {
2178 0
        ilow = nop-1;
2179
    }
2180 1
    if (ihigh < ilow) {
2181
        ihigh = ilow;
2182
    }
2183 1
    else if (ihigh > nop) {
2184 0
        ihigh = nop;
2185
    }
2186

2187 1
    if (!PySequence_Check(v) || PySequence_Size(v) != ihigh-ilow) {
2188 0
        PyErr_SetString(PyExc_ValueError,
2189
                "Wrong size to assign to iterator slice");
2190 0
        return -1;
2191
    }
2192

2193 1
    for (i = ilow; i < ihigh ; ++i) {
2194 1
        PyObject *item = PySequence_GetItem(v, i-ilow);
2195 1
        if (item == NULL) {
2196
            return -1;
2197
        }
2198 1
        if (npyiter_seq_ass_item(self, i, item) < 0) {
2199 0
            Py_DECREF(item);
2200
            return -1;
2201
        }
2202 1
        Py_DECREF(item);
2203
    }
2204

2205
    return 0;
2206
}
2207

2208
static PyObject *
2209 1
npyiter_subscript(NewNpyArrayIterObject *self, PyObject *op)
2210
{
2211 1
    if (self->iter == NULL || self->finished) {
2212 0
        PyErr_SetString(PyExc_ValueError,
2213
                "Iterator is past the end");
2214 0
        return NULL;
2215
    }
2216

2217 1
    if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2218 1
        PyErr_SetString(PyExc_ValueError,
2219
                "Iterator construction used delayed buffer allocation, "
2220
                "and no reset has been done yet");
2221 1
        return NULL;
2222
    }
2223

2224 1
    if (PyLong_Check(op) ||
2225 1
                    (PyIndex_Check(op) && !PySequence_Check(op))) {
2226 1
        npy_intp i = PyArray_PyIntAsIntp(op);
2227 1
        if (error_converting(i)) {
2228
            return NULL;
2229
        }
2230 1
        return npyiter_seq_item(self, i);
2231
    }
2232 1
    else if (PySlice_Check(op)) {
2233 1
        Py_ssize_t istart = 0, iend = 0, istep = 0, islicelength;
2234 1
        if (PySlice_GetIndicesEx(op, NpyIter_GetNOp(self->iter),
2235 1
                                 &istart, &iend, &istep, &islicelength) < 0) {
2236
            return NULL;
2237
        }
2238 1
        if (istep != 1) {
2239 0
            PyErr_SetString(PyExc_ValueError,
2240
                    "Iterator slicing only supports a step of 1");
2241 0
            return NULL;
2242
        }
2243 1
        return npyiter_seq_slice(self, istart, iend);
2244
    }
2245

2246 0
    PyErr_SetString(PyExc_TypeError,
2247
            "invalid index type for iterator indexing");
2248 0
    return NULL;
2249
}
2250

2251
static int
2252 1
npyiter_ass_subscript(NewNpyArrayIterObject *self, PyObject *op,
2253
                        PyObject *value)
2254
{
2255 1
    if (value == NULL) {
2256 1
        PyErr_SetString(PyExc_TypeError,
2257
                "Cannot delete iterator elements");
2258 1
        return -1;
2259
    }
2260 1
    if (self->iter == NULL || self->finished) {
2261 0
        PyErr_SetString(PyExc_ValueError,
2262
                "Iterator is past the end");
2263 0
        return -1;
2264
    }
2265

2266 1
    if (NpyIter_HasDelayedBufAlloc(self->iter)) {
2267 1
        PyErr_SetString(PyExc_ValueError,
2268
                "Iterator construction used delayed buffer allocation, "
2269
                "and no reset has been done yet");
2270 1
        return -1;
2271
    }
2272

2273 1
    if (PyLong_Check(op) ||
2274 1
                    (PyIndex_Check(op) && !PySequence_Check(op))) {
2275 1
        npy_intp i = PyArray_PyIntAsIntp(op);
2276 1
        if (error_converting(i)) {
2277
            return -1;
2278
        }
2279 1
        return npyiter_seq_ass_item(self, i, value);
2280
    }
2281 1
    else if (PySlice_Check(op)) {
2282 1
        Py_ssize_t istart = 0, iend = 0, istep = 0, islicelength = 0;
2283 1
        if (PySlice_GetIndicesEx(op, NpyIter_GetNOp(self->iter),
2284 1
                                 &istart, &iend, &istep, &islicelength) < 0) {
2285
            return -1;
2286
        }
2287 1
        if (istep != 1) {
2288 0
            PyErr_SetString(PyExc_ValueError,
2289
                    "Iterator slice assignment only supports a step of 1");
2290 0
            return -1;
2291
        }
2292 1
        return npyiter_seq_ass_slice(self, istart, iend, value);
2293
    }
2294

2295 0
    PyErr_SetString(PyExc_TypeError,
2296
            "invalid index type for iterator indexing");
2297 0
    return -1;
2298
}
2299

2300
static PyObject *
2301 1
npyiter_enter(NewNpyArrayIterObject *self)
2302
{
2303 1
    if (self->iter == NULL) {
2304 1
        PyErr_SetString(PyExc_RuntimeError, "operation on non-initialized iterator");
2305 1
        return NULL;
2306
    }
2307 1
    Py_INCREF(self);
2308 1
    return (PyObject *)self;
2309
}
2310

2311
static PyObject *
2312 1
npyiter_close(NewNpyArrayIterObject *self)
2313
{
2314 1
    NpyIter *iter = self->iter;
2315
    int ret;
2316 1
    if (self->iter == NULL) {
2317 1
        Py_RETURN_NONE;
2318
    }
2319 1
    ret = NpyIter_Deallocate(iter);
2320 1
    self->iter = NULL;
2321 1
    Py_XDECREF(self->nested_child);
2322 1
    self->nested_child = NULL;
2323 1
    if (ret < 0) {
2324
        return NULL;
2325
    }
2326 1
    Py_RETURN_NONE;
2327
}
2328

2329
static PyObject *
2330 1
npyiter_exit(NewNpyArrayIterObject *self, PyObject *NPY_UNUSED(args))
2331
{
2332
    /* even if called via exception handling, writeback any data */
2333 1
    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
};

Read our documentation on viewing source code .

Loading