1
/*
2
 * This file implements type resolution for NumPy element-wise ufuncs.
3
 * This mechanism is still backwards-compatible with the pre-existing
4
 * legacy mechanism, so performs much slower than is necessary.
5
 *
6
 * Written by Mark Wiebe (mwwiebe@gmail.com)
7
 * Copyright (c) 2011 by Enthought, Inc.
8
 *
9
 * See LICENSE.txt for the license.
10
 */
11
#define _UMATHMODULE
12
#define _MULTIARRAYMODULE
13
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
14

15
#include <stdbool.h>
16

17
#include "Python.h"
18

19
#include "npy_config.h"
20
#include "npy_pycompat.h"
21
#include "npy_import.h"
22

23
#include "numpy/ufuncobject.h"
24
#include "ufunc_type_resolution.h"
25
#include "ufunc_object.h"
26
#include "common.h"
27
#include "convert_datatype.h"
28

29
#include "mem_overlap.h"
30
#if defined(HAVE_CBLAS)
31
#include "cblasfuncs.h"
32
#endif
33

34
static PyObject *
35 1
npy_casting_to_py_object(NPY_CASTING casting)
36
{
37 1
    switch (casting) {
38 0
        case NPY_NO_CASTING:
39 0
            return PyUnicode_FromString("no");
40 0
        case NPY_EQUIV_CASTING:
41 0
            return PyUnicode_FromString("equiv");
42 0
        case NPY_SAFE_CASTING:
43 0
            return PyUnicode_FromString("safe");
44 1
        case NPY_SAME_KIND_CASTING:
45 1
            return PyUnicode_FromString("same_kind");
46 0
        case NPY_UNSAFE_CASTING:
47 0
            return PyUnicode_FromString("unsafe");
48 0
        default:
49 0
            return PyLong_FromLong(casting);
50
    }
51
}
52

53

54
/**
55
 * Always returns -1 to indicate the exception was raised, for convenience
56
 */
57
static int
58 1
raise_binary_type_reso_error(PyUFuncObject *ufunc, PyArrayObject **operands) {
59
    static PyObject *exc_type = NULL;
60
    PyObject *exc_value;
61

62 1
    npy_cache_import(
63
        "numpy.core._exceptions", "_UFuncBinaryResolutionError",
64
        &exc_type);
65 1
    if (exc_type == NULL) {
66
        return -1;
67
    }
68

69
    /* produce an error object */
70 1
    exc_value = Py_BuildValue(
71
        "O(OO)", ufunc,
72 1
        (PyObject *)PyArray_DESCR(operands[0]),
73 1
        (PyObject *)PyArray_DESCR(operands[1])
74
    );
75 1
    if (exc_value == NULL){
76
        return -1;
77
    }
78 1
    PyErr_SetObject(exc_type, exc_value);
79 1
    Py_DECREF(exc_value);
80

81
    return -1;
82
}
83

84
/** Helper function to raise UFuncNoLoopError
85
 * Always returns -1 to indicate the exception was raised, for convenience
86
 */
87
static int
88 1
raise_no_loop_found_error(
89
        PyUFuncObject *ufunc, PyArray_Descr **dtypes)
90
{
91
    static PyObject *exc_type = NULL;
92
    PyObject *exc_value;
93
    PyObject *dtypes_tup;
94
    npy_intp i;
95

96 1
    npy_cache_import(
97
        "numpy.core._exceptions", "_UFuncNoLoopError",
98
        &exc_type);
99 1
    if (exc_type == NULL) {
100
        return -1;
101
    }
102

103
    /* convert dtypes to a tuple */
104 1
    dtypes_tup = PyTuple_New(ufunc->nargs);
105 1
    if (dtypes_tup == NULL) {
106
        return -1;
107
    }
108 1
    for (i = 0; i < ufunc->nargs; ++i) {
109 1
        Py_INCREF(dtypes[i]);
110 1
        PyTuple_SET_ITEM(dtypes_tup, i, (PyObject *)dtypes[i]);
111
    }
112

113
    /* produce an error object */
114 1
    exc_value = PyTuple_Pack(2, ufunc, dtypes_tup);
115 1
    Py_DECREF(dtypes_tup);
116 1
    if (exc_value == NULL){
117
        return -1;
118
    }
119 1
    PyErr_SetObject(exc_type, exc_value);
120 1
    Py_DECREF(exc_value);
121

122
    return -1;
123
}
124

125
static int
126 1
raise_casting_error(
127
        PyObject *exc_type,
128
        PyUFuncObject *ufunc,
129
        NPY_CASTING casting,
130
        PyArray_Descr *from,
131
        PyArray_Descr *to,
132
        npy_intp i)
133
{
134
    PyObject *exc_value;
135
    PyObject *casting_value;
136

137 1
    casting_value = npy_casting_to_py_object(casting);
138 1
    if (casting_value == NULL) {
139
        return -1;
140
    }
141

142 1
    exc_value = Py_BuildValue(
143
        "ONOOi",
144
        ufunc,
145
        casting_value,
146
        (PyObject *)from,
147
        (PyObject *)to,
148
        i
149
    );
150 1
    if (exc_value == NULL){
151
        return -1;
152
    }
153 1
    PyErr_SetObject(exc_type, exc_value);
154 1
    Py_DECREF(exc_value);
155

156
    return -1;
157
}
158

159
/** Helper function to raise UFuncInputCastingError
160
 * Always returns -1 to indicate the exception was raised, for convenience
161
 */
162
static int
163 1
raise_input_casting_error(
164
        PyUFuncObject *ufunc,
165
        NPY_CASTING casting,
166
        PyArray_Descr *from,
167
        PyArray_Descr *to,
168
        npy_intp i)
169
{
170
    static PyObject *exc_type = NULL;
171 1
    npy_cache_import(
172
        "numpy.core._exceptions", "_UFuncInputCastingError",
173
        &exc_type);
174 1
    if (exc_type == NULL) {
175
        return -1;
176
    }
177

178 1
    return raise_casting_error(exc_type, ufunc, casting, from, to, i);
179
}
180

181

182
/** Helper function to raise UFuncOutputCastingError
183
 * Always returns -1 to indicate the exception was raised, for convenience
184
 */
185
static int
186 1
raise_output_casting_error(
187
        PyUFuncObject *ufunc,
188
        NPY_CASTING casting,
189
        PyArray_Descr *from,
190
        PyArray_Descr *to,
191
        npy_intp i)
192
{
193
    static PyObject *exc_type = NULL;
194 1
    npy_cache_import(
195
        "numpy.core._exceptions", "_UFuncOutputCastingError",
196
        &exc_type);
197 1
    if (exc_type == NULL) {
198
        return -1;
199
    }
200

201 1
    return raise_casting_error(exc_type, ufunc, casting, from, to, i);
202
}
203

204

205
/*UFUNC_API
206
 *
207
 * Validates that the input operands can be cast to
208
 * the input types, and the output types can be cast to
209
 * the output operands where provided.
210
 *
211
 * Returns 0 on success, -1 (with exception raised) on validation failure.
212
 */
213
NPY_NO_EXPORT int
214 1
PyUFunc_ValidateCasting(PyUFuncObject *ufunc,
215
                            NPY_CASTING casting,
216
                            PyArrayObject **operands,
217
                            PyArray_Descr **dtypes)
218
{
219 1
    int i, nin = ufunc->nin, nop = nin + ufunc->nout;
220

221 1
    for (i = 0; i < nop; ++i) {
222 1
        if (i < nin) {
223 1
            if (!PyArray_CanCastArrayTo(operands[i], dtypes[i], casting)) {
224 1
                return raise_input_casting_error(
225
                    ufunc, casting, PyArray_DESCR(operands[i]), dtypes[i], i);
226
            }
227 1
        } else if (operands[i] != NULL) {
228 1
            if (!PyArray_CanCastTypeTo(dtypes[i],
229
                                    PyArray_DESCR(operands[i]), casting)) {
230 1
                return raise_output_casting_error(
231
                    ufunc, casting, dtypes[i], PyArray_DESCR(operands[i]), i);
232
            }
233
        }
234
    }
235

236
    return 0;
237
}
238

239
/*
240
 * Returns a new reference to type if it is already NBO, otherwise
241
 * returns a copy converted to NBO.
242
 */
243
static PyArray_Descr *
244
ensure_dtype_nbo(PyArray_Descr *type)
245
{
246 1
    if (PyArray_ISNBO(type->byteorder)) {
247 1
        Py_INCREF(type);
248
        return type;
249
    }
250
    else {
251 1
        return PyArray_DescrNewByteorder(type, NPY_NATIVE);
252
    }
253
}
254

255
/*UFUNC_API
256
 *
257
 * This function applies the default type resolution rules
258
 * for the provided ufunc.
259
 *
260
 * Returns 0 on success, -1 on error.
261
 */
262
NPY_NO_EXPORT int
263 1
PyUFunc_DefaultTypeResolver(PyUFuncObject *ufunc,
264
                                NPY_CASTING casting,
265
                                PyArrayObject **operands,
266
                                PyObject *type_tup,
267
                                PyArray_Descr **out_dtypes)
268
{
269 1
    int i, nop = ufunc->nin + ufunc->nout;
270 1
    int retval = 0, any_object = 0;
271
    NPY_CASTING input_casting;
272

273 1
    for (i = 0; i < nop; ++i) {
274 1
        if (operands[i] != NULL &&
275 1
                PyTypeNum_ISOBJECT(PyArray_DESCR(operands[i])->type_num)) {
276
            any_object = 1;
277
            break;
278
        }
279
    }
280

281
    /*
282
     * Decide the casting rules for inputs and outputs.  We want
283
     * NPY_SAFE_CASTING or stricter, so that the loop selection code
284
     * doesn't choose an integer loop for float inputs, or a float32
285
     * loop for float64 inputs.
286
     */
287 1
    input_casting = (casting > NPY_SAFE_CASTING) ? NPY_SAFE_CASTING : casting;
288

289 1
    if (type_tup == NULL) {
290
        /* Find the best ufunc inner loop, and fill in the dtypes */
291 1
        retval = linear_search_type_resolver(ufunc, operands,
292
                        input_casting, casting, any_object,
293
                        out_dtypes);
294
    } else {
295
        /* Find the specified ufunc inner loop, and fill in the dtypes */
296 1
        retval = type_tuple_type_resolver(ufunc, type_tup,
297
                        operands, casting, any_object, out_dtypes);
298
    }
299

300 1
    return retval;
301
}
302

303
/*
304
 * This function applies special type resolution rules for the case
305
 * where all the functions have the pattern XX->bool, using
306
 * PyArray_ResultType instead of a linear search to get the best
307
 * loop.
308
 *
309
 * Returns 0 on success, -1 on error.
310
 */
311
NPY_NO_EXPORT int
312 1
PyUFunc_SimpleBinaryComparisonTypeResolver(PyUFuncObject *ufunc,
313
                                NPY_CASTING casting,
314
                                PyArrayObject **operands,
315
                                PyObject *type_tup,
316
                                PyArray_Descr **out_dtypes)
317
{
318
    int i, type_num1, type_num2;
319 1
    const char *ufunc_name = ufunc_get_name_cstr(ufunc);
320

321 1
    if (ufunc->nin != 2 || ufunc->nout != 1) {
322 0
        PyErr_Format(PyExc_RuntimeError, "ufunc %s is configured "
323
                "to use binary comparison type resolution but has "
324
                "the wrong number of inputs or outputs",
325
                ufunc_name);
326 0
        return -1;
327
    }
328

329
    /*
330
     * Use the default type resolution if there's a custom data type
331
     * or object arrays.
332
     */
333 1
    type_num1 = PyArray_DESCR(operands[0])->type_num;
334 1
    type_num2 = PyArray_DESCR(operands[1])->type_num;
335 1
    if (type_num1 >= NPY_NTYPES || type_num2 >= NPY_NTYPES ||
336 1
            type_num1 == NPY_OBJECT || type_num2 == NPY_OBJECT) {
337 1
        return PyUFunc_DefaultTypeResolver(ufunc, casting, operands,
338
                type_tup, out_dtypes);
339
    }
340

341 1
    if (type_tup == NULL) {
342
        /* Input types are the result type */
343 1
        out_dtypes[0] = PyArray_ResultType(2, operands, 0, NULL);
344 1
        if (out_dtypes[0] == NULL) {
345
            return -1;
346
        }
347 1
        out_dtypes[1] = out_dtypes[0];
348 1
        Py_INCREF(out_dtypes[1]);
349
    }
350
    else {
351
        PyObject *item;
352 1
        PyArray_Descr *dtype = NULL;
353

354
        /*
355
         * If the type tuple isn't a single-element tuple, let the
356
         * default type resolution handle this one.
357
         */
358 1
        if (!PyTuple_Check(type_tup) || PyTuple_GET_SIZE(type_tup) != 1) {
359 1
            return PyUFunc_DefaultTypeResolver(ufunc, casting,
360
                    operands, type_tup, out_dtypes);
361
        }
362

363 0
        item = PyTuple_GET_ITEM(type_tup, 0);
364

365 0
        if (item == Py_None) {
366 0
            PyErr_SetString(PyExc_ValueError,
367
                    "require data type in the type tuple");
368 0
            return -1;
369
        }
370 0
        else if (!PyArray_DescrConverter(item, &dtype)) {
371
            return -1;
372
        }
373

374 0
        out_dtypes[0] = ensure_dtype_nbo(dtype);
375 0
        if (out_dtypes[0] == NULL) {
376
            return -1;
377
        }
378 0
        out_dtypes[1] = out_dtypes[0];
379 0
        Py_INCREF(out_dtypes[1]);
380
    }
381

382
    /* Output type is always boolean */
383 1
    out_dtypes[2] = PyArray_DescrFromType(NPY_BOOL);
384 1
    if (out_dtypes[2] == NULL) {
385 0
        for (i = 0; i < 2; ++i) {
386 0
            Py_DECREF(out_dtypes[i]);
387 0
            out_dtypes[i] = NULL;
388
        }
389
        return -1;
390
    }
391

392
    /* Check against the casting rules */
393 1
    if (PyUFunc_ValidateCasting(ufunc, casting, operands, out_dtypes) < 0) {
394 1
        for (i = 0; i < 3; ++i) {
395 1
            Py_DECREF(out_dtypes[i]);
396 1
            out_dtypes[i] = NULL;
397
        }
398
        return -1;
399
    }
400

401
    return 0;
402
}
403

404
NPY_NO_EXPORT int
405 1
PyUFunc_NegativeTypeResolver(PyUFuncObject *ufunc,
406
                             NPY_CASTING casting,
407
                             PyArrayObject **operands,
408
                             PyObject *type_tup,
409
                             PyArray_Descr **out_dtypes)
410
{
411
    int ret;
412 1
    ret = PyUFunc_SimpleUniformOperationTypeResolver(ufunc, casting, operands,
413
                                                   type_tup, out_dtypes);
414 1
    if (ret < 0) {
415
        return ret;
416
    }
417

418
    /* The type resolver would have upcast already */
419 1
    if (out_dtypes[0]->type_num == NPY_BOOL) {
420 1
        PyErr_Format(PyExc_TypeError,
421
            "The numpy boolean negative, the `-` operator, is not supported, "
422
            "use the `~` operator or the logical_not function instead.");
423 1
        return -1;
424
    }
425

426
    return ret;
427
}
428

429

430
/*
431
 * The ones_like function shouldn't really be a ufunc, but while it
432
 * still is, this provides type resolution that always forces UNSAFE
433
 * casting.
434
 */
435
NPY_NO_EXPORT int
436 1
PyUFunc_OnesLikeTypeResolver(PyUFuncObject *ufunc,
437
                                NPY_CASTING NPY_UNUSED(casting),
438
                                PyArrayObject **operands,
439
                                PyObject *type_tup,
440
                                PyArray_Descr **out_dtypes)
441
{
442 1
    return PyUFunc_SimpleUniformOperationTypeResolver(ufunc,
443
                        NPY_UNSAFE_CASTING,
444
                        operands, type_tup, out_dtypes);
445
}
446

447
/*
448
 * This function applies special type resolution rules for the case
449
 * where all of the types in the signature are the same, eg XX->X or XX->XX.
450
 * It uses PyArray_ResultType instead of a linear search to get the best
451
 * loop.
452
 *
453
 * Note that a simpler linear search through the functions loop
454
 * is still done, but switching to a simple array lookup for
455
 * built-in types would be better at some point.
456
 *
457
 * Returns 0 on success, -1 on error.
458
 */
459
NPY_NO_EXPORT int
460 1
PyUFunc_SimpleUniformOperationTypeResolver(
461
        PyUFuncObject *ufunc,
462
        NPY_CASTING casting,
463
        PyArrayObject **operands,
464
        PyObject *type_tup,
465
        PyArray_Descr **out_dtypes)
466
{
467 1
    const char *ufunc_name = ufunc_get_name_cstr(ufunc);
468

469 1
    if (ufunc->nin < 1) {
470 0
        PyErr_Format(PyExc_RuntimeError, "ufunc %s is configured "
471
                "to use uniform operation type resolution but has "
472
                "no inputs",
473
                ufunc_name);
474 0
        return -1;
475
    }
476 1
    int nop = ufunc->nin + ufunc->nout;
477

478
    /*
479
     * There's a custom data type or an object array
480
     */
481 1
    bool has_custom_or_object = false;
482 1
    for (int iop = 0; iop < ufunc->nin; iop++) {
483 1
        int type_num = PyArray_DESCR(operands[iop])->type_num;
484 1
        if (type_num >= NPY_NTYPES || type_num == NPY_OBJECT) {
485
            has_custom_or_object = true;
486
            break;
487
        }
488
    }
489

490 1
    if (has_custom_or_object) {
491 1
        return PyUFunc_DefaultTypeResolver(ufunc, casting, operands,
492
                type_tup, out_dtypes);
493
    }
494

495 1
    if (type_tup == NULL) {
496
        /* PyArray_ResultType forgets to force a byte order when n == 1 */
497 1
        if (ufunc->nin == 1){
498 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
499
        }
500
        else {
501 1
            out_dtypes[0] = PyArray_ResultType(ufunc->nin, operands, 0, NULL);
502
        }
503 1
        if (out_dtypes[0] == NULL) {
504
            return -1;
505
        }
506
    }
507
    else {
508
        PyObject *item;
509 1
        PyArray_Descr *dtype = NULL;
510

511
        /*
512
         * If the type tuple isn't a single-element tuple, let the
513
         * default type resolution handle this one.
514
         */
515 1
        if (!PyTuple_Check(type_tup) || PyTuple_GET_SIZE(type_tup) != 1) {
516 1
            return PyUFunc_DefaultTypeResolver(ufunc, casting,
517
                    operands, type_tup, out_dtypes);
518
        }
519

520 1
        item = PyTuple_GET_ITEM(type_tup, 0);
521

522 1
        if (item == Py_None) {
523 0
            PyErr_SetString(PyExc_ValueError,
524
                    "require data type in the type tuple");
525 0
            return -1;
526
        }
527 1
        else if (!PyArray_DescrConverter(item, &dtype)) {
528
            return -1;
529
        }
530

531 1
        out_dtypes[0] = ensure_dtype_nbo(dtype);
532 1
        Py_DECREF(dtype);
533 1
        if (out_dtypes[0] == NULL) {
534
            return -1;
535
        }
536
    }
537

538
    /* All types are the same - copy the first one to the rest */
539 1
    for (int iop = 1; iop < nop; iop++) {
540 1
        out_dtypes[iop] = out_dtypes[0];
541 1
        Py_INCREF(out_dtypes[iop]);
542
    }
543

544
    /* Check against the casting rules */
545 1
    if (PyUFunc_ValidateCasting(ufunc, casting, operands, out_dtypes) < 0) {
546 1
        for (int iop = 0; iop < nop; iop++) {
547 1
            Py_DECREF(out_dtypes[iop]);
548 1
            out_dtypes[iop] = NULL;
549
        }
550
        return -1;
551
    }
552

553
    return 0;
554
}
555

556
/*
557
 * This function applies special type resolution rules for the absolute
558
 * ufunc. This ufunc converts complex -> float, so isn't covered
559
 * by the simple unary type resolution.
560
 *
561
 * Returns 0 on success, -1 on error.
562
 */
563
NPY_NO_EXPORT int
564 1
PyUFunc_AbsoluteTypeResolver(PyUFuncObject *ufunc,
565
                                NPY_CASTING casting,
566
                                PyArrayObject **operands,
567
                                PyObject *type_tup,
568
                                PyArray_Descr **out_dtypes)
569
{
570
    /* Use the default for complex types, to find the loop producing float */
571 1
    if (PyTypeNum_ISCOMPLEX(PyArray_DESCR(operands[0])->type_num)) {
572 1
        return PyUFunc_DefaultTypeResolver(ufunc, casting, operands,
573
                    type_tup, out_dtypes);
574
    }
575
    else {
576 1
        return PyUFunc_SimpleUniformOperationTypeResolver(ufunc, casting,
577
                    operands, type_tup, out_dtypes);
578
    }
579
}
580

581
/*
582
 * This function applies special type resolution rules for the isnat
583
 * ufunc. This ufunc converts datetime/timedelta -> bool, and is not covered
584
 * by the simple unary type resolution.
585
 *
586
 * Returns 0 on success, -1 on error.
587
 */
588
NPY_NO_EXPORT int
589 1
PyUFunc_IsNaTTypeResolver(PyUFuncObject *ufunc,
590
                          NPY_CASTING casting,
591
                          PyArrayObject **operands,
592
                          PyObject *type_tup,
593
                          PyArray_Descr **out_dtypes)
594
{
595 1
    if (!PyTypeNum_ISDATETIME(PyArray_DESCR(operands[0])->type_num)) {
596 1
        PyErr_SetString(PyExc_TypeError,
597
                "ufunc 'isnat' is only defined for datetime and timedelta.");
598 1
        return -1;
599
    }
600

601 1
    out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
602 1
    out_dtypes[1] = PyArray_DescrFromType(NPY_BOOL);
603

604 1
    return 0;
605
}
606

607

608
NPY_NO_EXPORT int
609 1
PyUFunc_IsFiniteTypeResolver(PyUFuncObject *ufunc,
610
                          NPY_CASTING casting,
611
                          PyArrayObject **operands,
612
                          PyObject *type_tup,
613
                          PyArray_Descr **out_dtypes)
614
{
615 1
    if (!PyTypeNum_ISDATETIME(PyArray_DESCR(operands[0])->type_num)) {
616 1
        return PyUFunc_DefaultTypeResolver(ufunc, casting, operands,
617
                                    type_tup, out_dtypes);
618
    }
619

620 1
    out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
621 1
    out_dtypes[1] = PyArray_DescrFromType(NPY_BOOL);
622

623 1
    return 0;
624
}
625

626

627
/*
628
 * Creates a new NPY_TIMEDELTA dtype, copying the datetime metadata
629
 * from the given dtype.
630
 *
631
 * NOTE: This function is copied from datetime.c in multiarray,
632
 *       because umath and multiarray are not linked together.
633
 */
634
static PyArray_Descr *
635
timedelta_dtype_with_copied_meta(PyArray_Descr *dtype)
636
{
637
    PyArray_Descr *ret;
638
    PyArray_DatetimeMetaData *dst, *src;
639
    PyArray_DatetimeDTypeMetaData *dst_dtmd, *src_dtmd;
640

641 1
    ret = PyArray_DescrNewFromType(NPY_TIMEDELTA);
642 1
    if (ret == NULL) {
643
        return NULL;
644
    }
645

646 1
    src_dtmd = ((PyArray_DatetimeDTypeMetaData *)dtype->c_metadata);
647 1
    dst_dtmd = ((PyArray_DatetimeDTypeMetaData *)ret->c_metadata);
648 1
    src = &(src_dtmd->meta);
649 1
    dst = &(dst_dtmd->meta);
650

651 1
    *dst = *src;
652

653
    return ret;
654
}
655

656
/*
657
 * This function applies the type resolution rules for addition.
658
 * In particular, there are a number of special cases with datetime:
659
 *    m8[<A>] + m8[<B>] => m8[gcd(<A>,<B>)] + m8[gcd(<A>,<B>)]
660
 *    m8[<A>] + int     => m8[<A>] + m8[<A>]
661
 *    int     + m8[<A>] => m8[<A>] + m8[<A>]
662
 *    M8[<A>] + int     => M8[<A>] + m8[<A>]
663
 *    int     + M8[<A>] => m8[<A>] + M8[<A>]
664
 *    M8[<A>] + m8[<B>] => M8[gcd(<A>,<B>)] + m8[gcd(<A>,<B>)]
665
 *    m8[<A>] + M8[<B>] => m8[gcd(<A>,<B>)] + M8[gcd(<A>,<B>)]
666
 * TODO: Non-linear time unit cases require highly special-cased loops
667
 *    M8[<A>] + m8[Y|M|B]
668
 *    m8[Y|M|B] + M8[<A>]
669
 */
670
NPY_NO_EXPORT int
671 1
PyUFunc_AdditionTypeResolver(PyUFuncObject *ufunc,
672
                                NPY_CASTING casting,
673
                                PyArrayObject **operands,
674
                                PyObject *type_tup,
675
                                PyArray_Descr **out_dtypes)
676
{
677
    int type_num1, type_num2;
678
    int i;
679

680 1
    type_num1 = PyArray_DESCR(operands[0])->type_num;
681 1
    type_num2 = PyArray_DESCR(operands[1])->type_num;
682

683
    /* Use the default when datetime and timedelta are not involved */
684 1
    if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
685 1
        return PyUFunc_SimpleUniformOperationTypeResolver(ufunc, casting,
686
                    operands, type_tup, out_dtypes);
687
    }
688

689 1
    if (type_num1 == NPY_TIMEDELTA) {
690
        /* m8[<A>] + m8[<B>] => m8[gcd(<A>,<B>)] + m8[gcd(<A>,<B>)] */
691 1
        if (type_num2 == NPY_TIMEDELTA) {
692 1
            out_dtypes[0] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
693
                                                PyArray_DESCR(operands[1]));
694 1
            if (out_dtypes[0] == NULL) {
695
                return -1;
696
            }
697 1
            out_dtypes[1] = out_dtypes[0];
698 1
            Py_INCREF(out_dtypes[1]);
699 1
            out_dtypes[2] = out_dtypes[0];
700 1
            Py_INCREF(out_dtypes[2]);
701
        }
702
        /* m8[<A>] + M8[<B>] => m8[gcd(<A>,<B>)] + M8[gcd(<A>,<B>)] */
703 1
        else if (type_num2 == NPY_DATETIME) {
704 1
            out_dtypes[1] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
705
                                                PyArray_DESCR(operands[1]));
706 1
            if (out_dtypes[1] == NULL) {
707
                return -1;
708
            }
709
            /* Make a new NPY_TIMEDELTA, and copy the datetime's metadata */
710 1
            out_dtypes[0] = timedelta_dtype_with_copied_meta(out_dtypes[1]);
711 1
            if (out_dtypes[0] == NULL) {
712 0
                Py_DECREF(out_dtypes[1]);
713 0
                out_dtypes[1] = NULL;
714 0
                return -1;
715
            }
716 1
            out_dtypes[2] = out_dtypes[1];
717 1
            Py_INCREF(out_dtypes[2]);
718
        }
719
        /* m8[<A>] + int => m8[<A>] + m8[<A>] */
720 1
        else if (PyTypeNum_ISINTEGER(type_num2) ||
721
                                    PyTypeNum_ISBOOL(type_num2)) {
722 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
723 1
            if (out_dtypes[0] == NULL) {
724
                return -1;
725
            }
726 1
            out_dtypes[1] = out_dtypes[0];
727 1
            Py_INCREF(out_dtypes[1]);
728 1
            out_dtypes[2] = out_dtypes[0];
729 1
            Py_INCREF(out_dtypes[2]);
730

731 1
            type_num2 = NPY_TIMEDELTA;
732
        }
733
        else {
734 0
            return raise_binary_type_reso_error(ufunc, operands);
735
        }
736
    }
737 1
    else if (type_num1 == NPY_DATETIME) {
738
        /* M8[<A>] + m8[<B>] => M8[gcd(<A>,<B>)] + m8[gcd(<A>,<B>)] */
739 1
        if (type_num2 == NPY_TIMEDELTA) {
740 1
            out_dtypes[0] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
741
                                                PyArray_DESCR(operands[1]));
742 1
            if (out_dtypes[0] == NULL) {
743
                return -1;
744
            }
745
            /* Make a new NPY_TIMEDELTA, and copy the datetime's metadata */
746 1
            out_dtypes[1] = timedelta_dtype_with_copied_meta(out_dtypes[0]);
747 1
            if (out_dtypes[1] == NULL) {
748 0
                Py_DECREF(out_dtypes[0]);
749 0
                out_dtypes[0] = NULL;
750 0
                return -1;
751
            }
752 1
            out_dtypes[2] = out_dtypes[0];
753 1
            Py_INCREF(out_dtypes[2]);
754
        }
755
        /* M8[<A>] + int => M8[<A>] + m8[<A>] */
756 1
        else if (PyTypeNum_ISINTEGER(type_num2) ||
757
                    PyTypeNum_ISBOOL(type_num2)) {
758 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
759 1
            if (out_dtypes[0] == NULL) {
760
                return -1;
761
            }
762
            /* Make a new NPY_TIMEDELTA, and copy type1's metadata */
763 1
            out_dtypes[1] = timedelta_dtype_with_copied_meta(
764
                                            PyArray_DESCR(operands[0]));
765 1
            if (out_dtypes[1] == NULL) {
766 0
                Py_DECREF(out_dtypes[0]);
767 0
                out_dtypes[0] = NULL;
768 0
                return -1;
769
            }
770 1
            out_dtypes[2] = out_dtypes[0];
771 1
            Py_INCREF(out_dtypes[2]);
772

773 1
            type_num2 = NPY_TIMEDELTA;
774
        }
775
        else {
776 1
            return raise_binary_type_reso_error(ufunc, operands);
777
        }
778
    }
779 1
    else if (PyTypeNum_ISINTEGER(type_num1) || PyTypeNum_ISBOOL(type_num1)) {
780
        /* int + m8[<A>] => m8[<A>] + m8[<A>] */
781 1
        if (type_num2 == NPY_TIMEDELTA) {
782 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[1]));
783 1
            if (out_dtypes[0] == NULL) {
784
                return -1;
785
            }
786 1
            out_dtypes[1] = out_dtypes[0];
787 1
            Py_INCREF(out_dtypes[1]);
788 1
            out_dtypes[2] = out_dtypes[0];
789 1
            Py_INCREF(out_dtypes[2]);
790

791 1
            type_num1 = NPY_TIMEDELTA;
792
        }
793 1
        else if (type_num2 == NPY_DATETIME) {
794
            /* Make a new NPY_TIMEDELTA, and copy type2's metadata */
795 1
            out_dtypes[0] = timedelta_dtype_with_copied_meta(
796
                                            PyArray_DESCR(operands[1]));
797 1
            if (out_dtypes[0] == NULL) {
798
                return -1;
799
            }
800 1
            out_dtypes[1] = ensure_dtype_nbo(PyArray_DESCR(operands[1]));
801 1
            if (out_dtypes[1] == NULL) {
802 0
                Py_DECREF(out_dtypes[0]);
803 0
                out_dtypes[0] = NULL;
804 0
                return -1;
805
            }
806 1
            out_dtypes[2] = out_dtypes[1];
807 1
            Py_INCREF(out_dtypes[2]);
808

809 1
            type_num1 = NPY_TIMEDELTA;
810
        }
811
        else {
812 0
            return raise_binary_type_reso_error(ufunc, operands);
813
        }
814
    }
815
    else {
816 0
        return raise_binary_type_reso_error(ufunc, operands);
817
    }
818

819
    /* Check against the casting rules */
820 1
    if (PyUFunc_ValidateCasting(ufunc, casting, operands, out_dtypes) < 0) {
821 0
        for (i = 0; i < 3; ++i) {
822 0
            Py_DECREF(out_dtypes[i]);
823 0
            out_dtypes[i] = NULL;
824
        }
825
        return -1;
826
    }
827

828
    return 0;
829
}
830

831
/*
832
 * This function applies the type resolution rules for subtraction.
833
 * In particular, there are a number of special cases with datetime:
834
 *    m8[<A>] - m8[<B>] => m8[gcd(<A>,<B>)] - m8[gcd(<A>,<B>)]
835
 *    m8[<A>] - int     => m8[<A>] - m8[<A>]
836
 *    int     - m8[<A>] => m8[<A>] - m8[<A>]
837
 *    M8[<A>] - int     => M8[<A>] - m8[<A>]
838
 *    M8[<A>] - m8[<B>] => M8[gcd(<A>,<B>)] - m8[gcd(<A>,<B>)]
839
 * TODO: Non-linear time unit cases require highly special-cased loops
840
 *    M8[<A>] - m8[Y|M|B]
841
 */
842
NPY_NO_EXPORT int
843 1
PyUFunc_SubtractionTypeResolver(PyUFuncObject *ufunc,
844
                                NPY_CASTING casting,
845
                                PyArrayObject **operands,
846
                                PyObject *type_tup,
847
                                PyArray_Descr **out_dtypes)
848
{
849
    int type_num1, type_num2;
850
    int i;
851

852 1
    type_num1 = PyArray_DESCR(operands[0])->type_num;
853 1
    type_num2 = PyArray_DESCR(operands[1])->type_num;
854

855
    /* Use the default when datetime and timedelta are not involved */
856 1
    if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
857
        int ret;
858 1
        ret = PyUFunc_SimpleUniformOperationTypeResolver(ufunc, casting,
859
                                                operands, type_tup, out_dtypes);
860 1
        if (ret < 0) {
861
            return ret;
862
        }
863

864
        /* The type resolver would have upcast already */
865 1
        if (out_dtypes[0]->type_num == NPY_BOOL) {
866 1
            PyErr_Format(PyExc_TypeError,
867
                "numpy boolean subtract, the `-` operator, is not supported, "
868
                "use the bitwise_xor, the `^` operator, or the logical_xor "
869
                "function instead.");
870 1
            return -1;
871
        }
872
        return ret;
873
    }
874

875 1
    if (type_num1 == NPY_TIMEDELTA) {
876
        /* m8[<A>] - m8[<B>] => m8[gcd(<A>,<B>)] - m8[gcd(<A>,<B>)] */
877 1
        if (type_num2 == NPY_TIMEDELTA) {
878 1
            out_dtypes[0] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
879
                                                PyArray_DESCR(operands[1]));
880 1
            if (out_dtypes[0] == NULL) {
881
                return -1;
882
            }
883 1
            out_dtypes[1] = out_dtypes[0];
884 1
            Py_INCREF(out_dtypes[1]);
885 1
            out_dtypes[2] = out_dtypes[0];
886 1
            Py_INCREF(out_dtypes[2]);
887
        }
888
        /* m8[<A>] - int => m8[<A>] - m8[<A>] */
889 1
        else if (PyTypeNum_ISINTEGER(type_num2) ||
890
                                        PyTypeNum_ISBOOL(type_num2)) {
891 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
892 1
            if (out_dtypes[0] == NULL) {
893
                return -1;
894
            }
895 1
            out_dtypes[1] = out_dtypes[0];
896 1
            Py_INCREF(out_dtypes[1]);
897 1
            out_dtypes[2] = out_dtypes[0];
898 1
            Py_INCREF(out_dtypes[2]);
899

900 1
            type_num2 = NPY_TIMEDELTA;
901
        }
902
        else {
903 1
            return raise_binary_type_reso_error(ufunc, operands);
904
        }
905
    }
906 1
    else if (type_num1 == NPY_DATETIME) {
907
        /* M8[<A>] - m8[<B>] => M8[gcd(<A>,<B>)] - m8[gcd(<A>,<B>)] */
908 1
        if (type_num2 == NPY_TIMEDELTA) {
909 1
            out_dtypes[0] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
910
                                                PyArray_DESCR(operands[1]));
911 1
            if (out_dtypes[0] == NULL) {
912
                return -1;
913
            }
914
            /* Make a new NPY_TIMEDELTA, and copy the datetime's metadata */
915 1
            out_dtypes[1] = timedelta_dtype_with_copied_meta(out_dtypes[0]);
916 1
            if (out_dtypes[1] == NULL) {
917 0
                Py_DECREF(out_dtypes[0]);
918 0
                out_dtypes[0] = NULL;
919 0
                return -1;
920
            }
921 1
            out_dtypes[2] = out_dtypes[0];
922 1
            Py_INCREF(out_dtypes[2]);
923
        }
924
        /* M8[<A>] - int => M8[<A>] - m8[<A>] */
925 1
        else if (PyTypeNum_ISINTEGER(type_num2) ||
926
                    PyTypeNum_ISBOOL(type_num2)) {
927 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
928 1
            if (out_dtypes[0] == NULL) {
929
                return -1;
930
            }
931
            /* Make a new NPY_TIMEDELTA, and copy type1's metadata */
932 1
            out_dtypes[1] = timedelta_dtype_with_copied_meta(
933
                                            PyArray_DESCR(operands[0]));
934 1
            if (out_dtypes[1] == NULL) {
935 0
                Py_DECREF(out_dtypes[0]);
936 0
                out_dtypes[0] = NULL;
937 0
                return -1;
938
            }
939 1
            out_dtypes[2] = out_dtypes[0];
940 1
            Py_INCREF(out_dtypes[2]);
941

942 1
            type_num2 = NPY_TIMEDELTA;
943
        }
944
        /* M8[<A>] - M8[<B>] => M8[gcd(<A>,<B>)] - M8[gcd(<A>,<B>)] */
945 1
        else if (type_num2 == NPY_DATETIME) {
946 1
            out_dtypes[0] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
947
                                                PyArray_DESCR(operands[1]));
948 1
            if (out_dtypes[0] == NULL) {
949
                return -1;
950
            }
951
            /* Make a new NPY_TIMEDELTA, and copy type1's metadata */
952 1
            out_dtypes[2] = timedelta_dtype_with_copied_meta(out_dtypes[0]);
953 1
            if (out_dtypes[2] == NULL) {
954 0
                Py_DECREF(out_dtypes[0]);
955
                return -1;
956
            }
957 1
            out_dtypes[1] = out_dtypes[0];
958 1
            Py_INCREF(out_dtypes[1]);
959
        }
960
        else {
961 0
            return raise_binary_type_reso_error(ufunc, operands);
962
        }
963
    }
964 1
    else if (PyTypeNum_ISINTEGER(type_num1) || PyTypeNum_ISBOOL(type_num1)) {
965
        /* int - m8[<A>] => m8[<A>] - m8[<A>] */
966 1
        if (type_num2 == NPY_TIMEDELTA) {
967 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[1]));
968 1
            if (out_dtypes[0] == NULL) {
969
                return -1;
970
            }
971 1
            out_dtypes[1] = out_dtypes[0];
972 1
            Py_INCREF(out_dtypes[1]);
973 1
            out_dtypes[2] = out_dtypes[0];
974 1
            Py_INCREF(out_dtypes[2]);
975

976 1
            type_num1 = NPY_TIMEDELTA;
977
        }
978
        else {
979 1
            return raise_binary_type_reso_error(ufunc, operands);
980
        }
981
    }
982
    else {
983 0
        return raise_binary_type_reso_error(ufunc, operands);
984
    }
985

986
    /* Check against the casting rules */
987 1
    if (PyUFunc_ValidateCasting(ufunc, casting, operands, out_dtypes) < 0) {
988 0
        for (i = 0; i < 3; ++i) {
989 0
            Py_DECREF(out_dtypes[i]);
990 0
            out_dtypes[i] = NULL;
991
        }
992
        return -1;
993
    }
994

995
    return 0;
996
}
997

998
/*
999
 * This function applies the type resolution rules for multiplication.
1000
 * In particular, there are a number of special cases with datetime:
1001
 *    int## * m8[<A>] => int64 * m8[<A>]
1002
 *    m8[<A>] * int## => m8[<A>] * int64
1003
 *    float## * m8[<A>] => float64 * m8[<A>]
1004
 *    m8[<A>] * float## => m8[<A>] * float64
1005
 */
1006
NPY_NO_EXPORT int
1007 1
PyUFunc_MultiplicationTypeResolver(PyUFuncObject *ufunc,
1008
                                NPY_CASTING casting,
1009
                                PyArrayObject **operands,
1010
                                PyObject *type_tup,
1011
                                PyArray_Descr **out_dtypes)
1012
{
1013
    int type_num1, type_num2;
1014
    int i;
1015

1016 1
    type_num1 = PyArray_DESCR(operands[0])->type_num;
1017 1
    type_num2 = PyArray_DESCR(operands[1])->type_num;
1018

1019
    /* Use the default when datetime and timedelta are not involved */
1020 1
    if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
1021 1
        return PyUFunc_SimpleUniformOperationTypeResolver(ufunc, casting,
1022
                    operands, type_tup, out_dtypes);
1023
    }
1024

1025 1
    if (type_num1 == NPY_TIMEDELTA) {
1026
        /* m8[<A>] * int## => m8[<A>] * int64 */
1027 1
        if (PyTypeNum_ISINTEGER(type_num2) || PyTypeNum_ISBOOL(type_num2)) {
1028 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
1029 1
            if (out_dtypes[0] == NULL) {
1030
                return -1;
1031
            }
1032 1
            out_dtypes[1] = PyArray_DescrNewFromType(NPY_LONGLONG);
1033 1
            if (out_dtypes[1] == NULL) {
1034 0
                Py_DECREF(out_dtypes[0]);
1035 0
                out_dtypes[0] = NULL;
1036 0
                return -1;
1037
            }
1038 1
            out_dtypes[2] = out_dtypes[0];
1039 1
            Py_INCREF(out_dtypes[2]);
1040

1041 1
            type_num2 = NPY_LONGLONG;
1042
        }
1043
        /* m8[<A>] * float## => m8[<A>] * float64 */
1044 1
        else if (PyTypeNum_ISFLOAT(type_num2)) {
1045 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
1046 1
            if (out_dtypes[0] == NULL) {
1047
                return -1;
1048
            }
1049 1
            out_dtypes[1] = PyArray_DescrNewFromType(NPY_DOUBLE);
1050 1
            if (out_dtypes[1] == NULL) {
1051 0
                Py_DECREF(out_dtypes[0]);
1052 0
                out_dtypes[0] = NULL;
1053 0
                return -1;
1054
            }
1055 1
            out_dtypes[2] = out_dtypes[0];
1056 1
            Py_INCREF(out_dtypes[2]);
1057

1058 1
            type_num2 = NPY_DOUBLE;
1059
        }
1060
        else {
1061 1
            return raise_binary_type_reso_error(ufunc, operands);
1062
        }
1063
    }
1064 1
    else if (PyTypeNum_ISINTEGER(type_num1) || PyTypeNum_ISBOOL(type_num1)) {
1065
        /* int## * m8[<A>] => int64 * m8[<A>] */
1066 1
        if (type_num2 == NPY_TIMEDELTA) {
1067 1
            out_dtypes[0] = PyArray_DescrNewFromType(NPY_LONGLONG);
1068 1
            if (out_dtypes[0] == NULL) {
1069
                return -1;
1070
            }
1071 1
            out_dtypes[1] = ensure_dtype_nbo(PyArray_DESCR(operands[1]));
1072 1
            if (out_dtypes[1] == NULL) {
1073 0
                Py_DECREF(out_dtypes[0]);
1074 0
                out_dtypes[0] = NULL;
1075 0
                return -1;
1076
            }
1077 1
            out_dtypes[2] = out_dtypes[1];
1078 1
            Py_INCREF(out_dtypes[2]);
1079

1080 1
            type_num1 = NPY_LONGLONG;
1081
        }
1082
        else {
1083 1
            return raise_binary_type_reso_error(ufunc, operands);
1084
        }
1085
    }
1086 1
    else if (PyTypeNum_ISFLOAT(type_num1)) {
1087
        /* float## * m8[<A>] => float64 * m8[<A>] */
1088 1
        if (type_num2 == NPY_TIMEDELTA) {
1089 1
            out_dtypes[0] = PyArray_DescrNewFromType(NPY_DOUBLE);
1090 1
            if (out_dtypes[0] == NULL) {
1091
                return -1;
1092
            }
1093 1
            out_dtypes[1] = ensure_dtype_nbo(PyArray_DESCR(operands[1]));
1094 1
            if (out_dtypes[1] == NULL) {
1095 0
                Py_DECREF(out_dtypes[0]);
1096 0
                out_dtypes[0] = NULL;
1097 0
                return -1;
1098
            }
1099 1
            out_dtypes[2] = out_dtypes[1];
1100 1
            Py_INCREF(out_dtypes[2]);
1101

1102 1
            type_num1 = NPY_DOUBLE;
1103
        }
1104
        else {
1105 1
            return raise_binary_type_reso_error(ufunc, operands);
1106
        }
1107
    }
1108
    else {
1109 1
        return raise_binary_type_reso_error(ufunc, operands);
1110
    }
1111

1112
    /* Check against the casting rules */
1113 1
    if (PyUFunc_ValidateCasting(ufunc, casting, operands, out_dtypes) < 0) {
1114 0
        for (i = 0; i < 3; ++i) {
1115 0
            Py_DECREF(out_dtypes[i]);
1116 0
            out_dtypes[i] = NULL;
1117
        }
1118
        return -1;
1119
    }
1120

1121
    return 0;
1122
}
1123

1124

1125
/*
1126
 * This function applies the type resolution rules for division.
1127
 * In particular, there are a number of special cases with datetime:
1128
 *    m8[<A>] / m8[<B>] to  m8[gcd(<A>,<B>)] / m8[gcd(<A>,<B>)]  -> float64
1129
 *    m8[<A>] / int##   to m8[<A>] / int64 -> m8[<A>]
1130
 *    m8[<A>] / float## to m8[<A>] / float64 -> m8[<A>]
1131
 */
1132
NPY_NO_EXPORT int
1133 1
PyUFunc_DivisionTypeResolver(PyUFuncObject *ufunc,
1134
                                NPY_CASTING casting,
1135
                                PyArrayObject **operands,
1136
                                PyObject *type_tup,
1137
                                PyArray_Descr **out_dtypes)
1138
{
1139
    int type_num1, type_num2;
1140
    int i;
1141

1142 1
    type_num1 = PyArray_DESCR(operands[0])->type_num;
1143 1
    type_num2 = PyArray_DESCR(operands[1])->type_num;
1144

1145
    /* Use the default when datetime and timedelta are not involved */
1146 1
    if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
1147 1
        return PyUFunc_DefaultTypeResolver(ufunc, casting, operands,
1148
                    type_tup, out_dtypes);
1149
    }
1150

1151 1
    if (type_num1 == NPY_TIMEDELTA) {
1152
        /*
1153
         * m8[<A>] / m8[<B>] to
1154
         * m8[gcd(<A>,<B>)] / m8[gcd(<A>,<B>)]  -> float64
1155
         */
1156 1
        if (type_num2 == NPY_TIMEDELTA) {
1157 1
            out_dtypes[0] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
1158
                                                PyArray_DESCR(operands[1]));
1159 1
            if (out_dtypes[0] == NULL) {
1160
                return -1;
1161
            }
1162 1
            out_dtypes[1] = out_dtypes[0];
1163 1
            Py_INCREF(out_dtypes[1]);
1164

1165
            /*
1166
             * TODO: split function into truediv and floordiv resolvers
1167
             */
1168 1
            if (strcmp(ufunc->name, "floor_divide") == 0) {
1169 1
                out_dtypes[2] = PyArray_DescrFromType(NPY_LONGLONG);
1170
            }
1171
            else {
1172 1
            out_dtypes[2] = PyArray_DescrFromType(NPY_DOUBLE);
1173
            }
1174 1
            if (out_dtypes[2] == NULL) {
1175 0
                Py_DECREF(out_dtypes[0]);
1176 0
                out_dtypes[0] = NULL;
1177 0
                Py_DECREF(out_dtypes[1]);
1178 0
                out_dtypes[1] = NULL;
1179 0
                return -1;
1180
            }
1181
        }
1182
        /* m8[<A>] / int## => m8[<A>] / int64 */
1183 1
        else if (PyTypeNum_ISINTEGER(type_num2)) {
1184 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
1185 1
            if (out_dtypes[0] == NULL) {
1186
                return -1;
1187
            }
1188 1
            out_dtypes[1] = PyArray_DescrFromType(NPY_LONGLONG);
1189 1
            if (out_dtypes[1] == NULL) {
1190 0
                Py_DECREF(out_dtypes[0]);
1191 0
                out_dtypes[0] = NULL;
1192 0
                return -1;
1193
            }
1194 1
            out_dtypes[2] = out_dtypes[0];
1195 1
            Py_INCREF(out_dtypes[2]);
1196

1197 1
            type_num2 = NPY_LONGLONG;
1198
        }
1199
        /* m8[<A>] / float## => m8[<A>] / float64 */
1200 1
        else if (PyTypeNum_ISFLOAT(type_num2)) {
1201 1
            out_dtypes[0] = ensure_dtype_nbo(PyArray_DESCR(operands[0]));
1202 1
            if (out_dtypes[0] == NULL) {
1203
                return -1;
1204
            }
1205 1
            out_dtypes[1] = PyArray_DescrNewFromType(NPY_DOUBLE);
1206 1
            if (out_dtypes[1] == NULL) {
1207 0
                Py_DECREF(out_dtypes[0]);
1208 0
                out_dtypes[0] = NULL;
1209 0
                return -1;
1210
            }
1211 1
            out_dtypes[2] = out_dtypes[0];
1212 1
            Py_INCREF(out_dtypes[2]);
1213

1214 1
            type_num2 = NPY_DOUBLE;
1215
        }
1216
        else {
1217 1
            return raise_binary_type_reso_error(ufunc, operands);
1218
        }
1219
    }
1220
    else {
1221 1
        return raise_binary_type_reso_error(ufunc, operands);
1222
    }
1223

1224
    /* Check against the casting rules */
1225 1
    if (PyUFunc_ValidateCasting(ufunc, casting, operands, out_dtypes) < 0) {
1226 0
        for (i = 0; i < 3; ++i) {
1227 0
            Py_DECREF(out_dtypes[i]);
1228 0
            out_dtypes[i] = NULL;
1229
        }
1230
        return -1;
1231
    }
1232

1233
    return 0;
1234
}
1235

1236

1237
NPY_NO_EXPORT int
1238 1
PyUFunc_RemainderTypeResolver(PyUFuncObject *ufunc,
1239
                                NPY_CASTING casting,
1240
                                PyArrayObject **operands,
1241
                                PyObject *type_tup,
1242
                                PyArray_Descr **out_dtypes)
1243
{
1244
    int type_num1, type_num2;
1245
    int i;
1246

1247 1
    type_num1 = PyArray_DESCR(operands[0])->type_num;
1248 1
    type_num2 = PyArray_DESCR(operands[1])->type_num;
1249

1250
    /* Use the default when datetime and timedelta are not involved */
1251 1
    if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
1252 1
        return PyUFunc_DefaultTypeResolver(ufunc, casting, operands,
1253
                    type_tup, out_dtypes);
1254
    }
1255 1
    if (type_num1 == NPY_TIMEDELTA) {
1256 1
        if (type_num2 == NPY_TIMEDELTA) {
1257 1
            out_dtypes[0] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
1258
                                                PyArray_DESCR(operands[1]));
1259 1
            if (out_dtypes[0] == NULL) {
1260
                return -1;
1261
            }
1262 1
            out_dtypes[1] = out_dtypes[0];
1263 1
            Py_INCREF(out_dtypes[1]);
1264 1
            out_dtypes[2] = out_dtypes[0];
1265 1
            Py_INCREF(out_dtypes[2]);
1266
        }
1267
        else {
1268 1
            return raise_binary_type_reso_error(ufunc, operands);
1269
        }
1270
    }
1271
    else {
1272 1
        return raise_binary_type_reso_error(ufunc, operands);
1273
    }
1274

1275
    /* Check against the casting rules */
1276 1
    if (PyUFunc_ValidateCasting(ufunc, casting, operands, out_dtypes) < 0) {
1277 0
        for (i = 0; i < 3; ++i) {
1278 0
            Py_DECREF(out_dtypes[i]);
1279 0
            out_dtypes[i] = NULL;
1280
        }
1281
        return -1;
1282
    }
1283

1284
    return 0;
1285
}
1286

1287

1288
/*
1289
 * True division should return float64 results when both inputs are integer
1290
 * types. The PyUFunc_DefaultTypeResolver promotes 8 bit integers to float16
1291
 * and 16 bit integers to float32, so that is overridden here by specifying a
1292
 * 'dd->d' signature. Returns -1 on failure.
1293
*/
1294
NPY_NO_EXPORT int
1295 1
PyUFunc_TrueDivisionTypeResolver(PyUFuncObject *ufunc,
1296
                                 NPY_CASTING casting,
1297
                                 PyArrayObject **operands,
1298
                                 PyObject *type_tup,
1299
                                 PyArray_Descr **out_dtypes)
1300
{
1301
    int type_num1, type_num2;
1302
    static PyObject *default_type_tup = NULL;
1303

1304
    /* Set default type for integer inputs to NPY_DOUBLE */
1305 1
    if (default_type_tup == NULL) {
1306 1
        PyArray_Descr *tmp = PyArray_DescrFromType(NPY_DOUBLE);
1307

1308 1
        if (tmp == NULL) {
1309
            return -1;
1310
        }
1311 1
        default_type_tup = PyTuple_Pack(3, tmp, tmp, tmp);
1312 1
        if (default_type_tup == NULL) {
1313 0
            Py_DECREF(tmp);
1314
            return -1;
1315
        }
1316 1
        Py_DECREF(tmp);
1317
    }
1318

1319 1
    type_num1 = PyArray_DESCR(operands[0])->type_num;
1320 1
    type_num2 = PyArray_DESCR(operands[1])->type_num;
1321

1322 1
    if (type_tup == NULL &&
1323 1
            (PyTypeNum_ISINTEGER(type_num1) || PyTypeNum_ISBOOL(type_num1)) &&
1324 1
            (PyTypeNum_ISINTEGER(type_num2) || PyTypeNum_ISBOOL(type_num2))) {
1325 1
        return PyUFunc_DefaultTypeResolver(ufunc, casting, operands,
1326
                                           default_type_tup, out_dtypes);
1327
    }
1328 1
    return PyUFunc_DivisionTypeResolver(ufunc, casting, operands,
1329
                                        type_tup, out_dtypes);
1330
}
1331

1332
static int
1333 1
find_userloop(PyUFuncObject *ufunc,
1334
                PyArray_Descr **dtypes,
1335
                PyUFuncGenericFunction *out_innerloop,
1336
                void **out_innerloopdata)
1337
{
1338 1
    npy_intp i, nin = ufunc->nin, j, nargs = nin + ufunc->nout;
1339

1340
    /* Use this to try to avoid repeating the same userdef loop search */
1341 1
    int last_userdef = -1;
1342

1343 1
    for (i = 0; i < nargs; ++i) {
1344
        int type_num;
1345

1346
        /* no more ufunc arguments to check */
1347 1
        if (dtypes[i] == NULL) {
1348
            break;
1349
        }
1350

1351 1
        type_num = dtypes[i]->type_num;
1352 1
        if (type_num != last_userdef &&
1353 1
                (PyTypeNum_ISUSERDEF(type_num) || type_num == NPY_VOID)) {
1354
            PyObject *key, *obj;
1355

1356 1
            last_userdef = type_num;
1357

1358 1
            key = PyLong_FromLong(type_num);
1359 1
            if (key == NULL) {
1360
                return -1;
1361
            }
1362 1
            obj = PyDict_GetItemWithError(ufunc->userloops, key);
1363 1
            Py_DECREF(key);
1364 1
            if (obj == NULL && PyErr_Occurred()){
1365
                return -1;
1366
            }
1367 1
            else if (obj == NULL) {
1368 0
                continue;
1369
            }
1370 1
            PyUFunc_Loop1d *funcdata = PyCapsule_GetPointer(obj, NULL);
1371 1
            if (funcdata == NULL) {
1372
                return -1;
1373
            }
1374 0
            for (; funcdata != NULL; funcdata = funcdata->next) {
1375 1
                int *types = funcdata->arg_types;
1376

1377 1
                for (j = 0; j < nargs; ++j) {
1378 1
                    if (types[j] != dtypes[j]->type_num) {
1379
                        break;
1380
                    }
1381
                }
1382
                /* It matched */
1383 1
                if (j == nargs) {
1384 1
                    *out_innerloop = funcdata->func;
1385 1
                    *out_innerloopdata = funcdata->data;
1386
                    return 1;
1387
                }
1388
            }
1389
        }
1390
    }
1391

1392
    /* Didn't find a match */
1393
    return 0;
1394
}
1395

1396
NPY_NO_EXPORT int
1397 1
PyUFunc_DefaultLegacyInnerLoopSelector(PyUFuncObject *ufunc,
1398
                                PyArray_Descr **dtypes,
1399
                                PyUFuncGenericFunction *out_innerloop,
1400
                                void **out_innerloopdata,
1401
                                int *out_needs_api)
1402
{
1403 1
    int nargs = ufunc->nargs;
1404
    char *types;
1405
    int i, j;
1406

1407
    /*
1408
     * If there are user-loops search them first.
1409
     * TODO: There needs to be a loop selection acceleration structure,
1410
     *       like a hash table.
1411
     */
1412 1
    if (ufunc->userloops) {
1413 1
        switch (find_userloop(ufunc, dtypes,
1414
                    out_innerloop, out_innerloopdata)) {
1415
            /* Error */
1416
            case -1:
1417
                return -1;
1418
            /* Found a loop */
1419 1
            case 1:
1420 1
                return 0;
1421
        }
1422
    }
1423

1424 1
    types = ufunc->types;
1425 1
    for (i = 0; i < ufunc->ntypes; ++i) {
1426
        /* Copy the types into an int array for matching */
1427 1
        for (j = 0; j < nargs; ++j) {
1428 1
            if (types[j] != dtypes[j]->type_num) {
1429
                break;
1430
            }
1431
        }
1432 1
        if (j == nargs) {
1433 1
            *out_innerloop = ufunc->functions[i];
1434 1
            *out_innerloopdata = ufunc->data[i];
1435 1
            return 0;
1436
        }
1437

1438 1
        types += nargs;
1439
    }
1440

1441 1
    return raise_no_loop_found_error(ufunc, dtypes);
1442
}
1443

1444
typedef struct {
1445
    NpyAuxData base;
1446
    PyUFuncGenericFunction unmasked_innerloop;
1447
    void *unmasked_innerloopdata;
1448
    int nargs;
1449
} _ufunc_masker_data;
1450

1451
static NpyAuxData *
1452 0
ufunc_masker_data_clone(NpyAuxData *data)
1453
{
1454
    _ufunc_masker_data *n;
1455

1456
    /* Allocate a new one */
1457 0
    n = (_ufunc_masker_data *)PyArray_malloc(sizeof(_ufunc_masker_data));
1458 0
    if (n == NULL) {
1459
        return NULL;
1460
    }
1461

1462
    /* Copy the data (unmasked data doesn't have object semantics) */
1463 0
    memcpy(n, data, sizeof(_ufunc_masker_data));
1464

1465 0
    return (NpyAuxData *)n;
1466
}
1467

1468
/*
1469
 * This function wraps a regular unmasked ufunc inner loop as a
1470
 * masked ufunc inner loop, only calling the function for
1471
 * elements where the mask is True.
1472
 */
1473
static void
1474 1
unmasked_ufunc_loop_as_masked(
1475
             char **dataptrs, npy_intp *strides,
1476
             char *mask, npy_intp mask_stride,
1477
             npy_intp loopsize,
1478
             NpyAuxData *innerloopdata)
1479
{
1480
    _ufunc_masker_data *data;
1481
    int iargs, nargs;
1482
    PyUFuncGenericFunction unmasked_innerloop;
1483
    void *unmasked_innerloopdata;
1484
    npy_intp subloopsize;
1485

1486
    /* Put the aux data into local variables */
1487 1
    data = (_ufunc_masker_data *)innerloopdata;
1488 1
    unmasked_innerloop = data->unmasked_innerloop;
1489 1
    unmasked_innerloopdata = data->unmasked_innerloopdata;
1490 1
    nargs = data->nargs;
1491

1492
    /* Process the data as runs of unmasked values */
1493
    do {
1494
        /* Skip masked values */
1495 1
        mask = npy_memchr(mask, 0, mask_stride, loopsize, &subloopsize, 1);
1496 1
        for (iargs = 0; iargs < nargs; ++iargs) {
1497 1
            dataptrs[iargs] += subloopsize * strides[iargs];
1498
        }
1499 1
        loopsize -= subloopsize;
1500
        /*
1501
         * Process unmasked values (assumes unmasked loop doesn't
1502
         * mess with the 'args' pointer values)
1503
         */
1504 1
        mask = npy_memchr(mask, 0, mask_stride, loopsize, &subloopsize, 0);
1505 1
        unmasked_innerloop(dataptrs, &subloopsize, strides,
1506
                                        unmasked_innerloopdata);
1507 1
        for (iargs = 0; iargs < nargs; ++iargs) {
1508 1
            dataptrs[iargs] += subloopsize * strides[iargs];
1509
        }
1510 1
        loopsize -= subloopsize;
1511 1
    } while (loopsize > 0);
1512
}
1513

1514

1515
/*
1516
 * This function wraps a legacy inner loop so it becomes masked.
1517
 *
1518
 * Returns 0 on success, -1 on error.
1519
 */
1520
NPY_NO_EXPORT int
1521 1
PyUFunc_DefaultMaskedInnerLoopSelector(PyUFuncObject *ufunc,
1522
                            PyArray_Descr **dtypes,
1523
                            PyArray_Descr *mask_dtype,
1524
                            npy_intp *NPY_UNUSED(fixed_strides),
1525
                            npy_intp NPY_UNUSED(fixed_mask_stride),
1526
                            PyUFunc_MaskedStridedInnerLoopFunc **out_innerloop,
1527
                            NpyAuxData **out_innerloopdata,
1528
                            int *out_needs_api)
1529
{
1530
    int retcode;
1531
    _ufunc_masker_data *data;
1532

1533 1
    if (ufunc->legacy_inner_loop_selector == NULL) {
1534 0
        PyErr_SetString(PyExc_RuntimeError,
1535
                "the ufunc default masked inner loop selector doesn't "
1536
                "yet support wrapping the new inner loop selector, it "
1537
                "still only wraps the legacy inner loop selector");
1538 0
        return -1;
1539
    }
1540

1541 1
    if (mask_dtype->type_num != NPY_BOOL) {
1542 0
        PyErr_SetString(PyExc_ValueError,
1543
                "only boolean masks are supported in ufunc inner loops "
1544
                "presently");
1545 0
        return -1;
1546
    }
1547

1548
    /* Create a new NpyAuxData object for the masker data */
1549 1
    data = (_ufunc_masker_data *)PyArray_malloc(sizeof(_ufunc_masker_data));
1550 1
    if (data == NULL) {
1551 0
        PyErr_NoMemory();
1552 0
        return -1;
1553
    }
1554 1
    memset(data, 0, sizeof(_ufunc_masker_data));
1555 1
    data->base.free = (NpyAuxData_FreeFunc *)&PyArray_free;
1556 1
    data->base.clone = &ufunc_masker_data_clone;
1557 1
    data->nargs = ufunc->nin + ufunc->nout;
1558

1559
    /* Get the unmasked ufunc inner loop */
1560 1
    retcode = ufunc->legacy_inner_loop_selector(ufunc, dtypes,
1561
                    &data->unmasked_innerloop, &data->unmasked_innerloopdata,
1562
                    out_needs_api);
1563 1
    if (retcode < 0) {
1564 0
        PyArray_free(data);
1565 0
        return retcode;
1566
    }
1567

1568
    /* Return the loop function + aux data */
1569 1
    *out_innerloop = &unmasked_ufunc_loop_as_masked;
1570 1
    *out_innerloopdata = (NpyAuxData *)data;
1571 1
    return 0;
1572
}
1573

1574
static int
1575 1
ufunc_loop_matches(PyUFuncObject *self,
1576
                    PyArrayObject **op,
1577
                    NPY_CASTING input_casting,
1578
                    NPY_CASTING output_casting,
1579
                    int any_object,
1580
                    int use_min_scalar,
1581
                    int *types, PyArray_Descr **dtypes,
1582
                    int *out_no_castable_output,
1583
                    char *out_err_src_typecode,
1584
                    char *out_err_dst_typecode)
1585
{
1586 1
    npy_intp i, nin = self->nin, nop = nin + self->nout;
1587

1588
    /*
1589
     * First check if all the inputs can be safely cast
1590
     * to the types for this function
1591
     */
1592 1
    for (i = 0; i < nin; ++i) {
1593
        PyArray_Descr *tmp;
1594

1595
        /*
1596
         * If no inputs are objects and there are more than one
1597
         * loop, don't allow conversion to object.  The rationale
1598
         * behind this is mostly performance.  Except for custom
1599
         * ufuncs built with just one object-parametered inner loop,
1600
         * only the types that are supported are implemented.  Trying
1601
         * the object version of logical_or on float arguments doesn't
1602
         * seem right.
1603
         */
1604 1
        if (types[i] == NPY_OBJECT && !any_object && self->ntypes > 1) {
1605
            return 0;
1606
        }
1607

1608
        /*
1609
         * If type num is NPY_VOID and struct dtypes have been passed in,
1610
         * use struct dtype object. Otherwise create new dtype object
1611
         * from type num.
1612
         */
1613 1
        if (types[i] == NPY_VOID && dtypes != NULL) {
1614 1
            tmp = dtypes[i];
1615 1
            Py_INCREF(tmp);
1616
        }
1617
        else {
1618 1
            tmp = PyArray_DescrFromType(types[i]);
1619
        }
1620 1
        if (tmp == NULL) {
1621
            return -1;
1622
        }
1623

1624
#if NPY_UF_DBG_TRACING
1625
        printf("Checking type for op %d, type %d: ", (int)i, (int)types[i]);
1626
        PyObject_Print((PyObject *)tmp, stdout, 0);
1627
        printf(", operand type: ");
1628
        PyObject_Print((PyObject *)PyArray_DESCR(op[i]), stdout, 0);
1629
        printf("\n");
1630
#endif
1631
        /*
1632
         * If all the inputs are scalars, use the regular
1633
         * promotion rules, not the special value-checking ones.
1634
         */
1635 1
        if (!use_min_scalar) {
1636 1
            if (!PyArray_CanCastTypeTo(PyArray_DESCR(op[i]), tmp,
1637
                                                    input_casting)) {
1638 1
                Py_DECREF(tmp);
1639
                return 0;
1640
            }
1641
        }
1642
        else {
1643 1
            if (!PyArray_CanCastArrayTo(op[i], tmp, input_casting)) {
1644 1
                Py_DECREF(tmp);
1645
                return 0;
1646
            }
1647
        }
1648 1
        Py_DECREF(tmp);
1649
    }
1650

1651
    /*
1652
     * If all the inputs were ok, then check casting back to the
1653
     * outputs.
1654
     */
1655 1
    for (i = nin; i < nop; ++i) {
1656 1
        if (op[i] != NULL) {
1657 1
            PyArray_Descr *tmp = PyArray_DescrFromType(types[i]);
1658 1
            if (tmp == NULL) {
1659
                return -1;
1660
            }
1661 1
            if (!PyArray_CanCastTypeTo(tmp, PyArray_DESCR(op[i]),
1662
                                                        output_casting)) {
1663 1
                if (!(*out_no_castable_output)) {
1664 1
                    *out_no_castable_output = 1;
1665 1
                    *out_err_src_typecode = tmp->type;
1666 1
                    *out_err_dst_typecode = PyArray_DESCR(op[i])->type;
1667
                }
1668 1
                Py_DECREF(tmp);
1669
                return 0;
1670
            }
1671 1
            Py_DECREF(tmp);
1672
        }
1673
    }
1674

1675
    return 1;
1676
}
1677

1678
static int
1679 1
set_ufunc_loop_data_types(PyUFuncObject *self, PyArrayObject **op,
1680
                    PyArray_Descr **out_dtypes,
1681
                    int *type_nums, PyArray_Descr **dtypes)
1682
{
1683 1
    int i, nin = self->nin, nop = nin + self->nout;
1684

1685
    /*
1686
     * Fill the dtypes array.
1687
     * For outputs,
1688
     * also search the inputs for a matching type_num to copy
1689
     * instead of creating a new one, similarly to preserve metadata.
1690
     **/
1691 1
    for (i = 0; i < nop; ++i) {
1692 1
        if (dtypes != NULL) {
1693 1
            out_dtypes[i] = dtypes[i];
1694 1
            Py_XINCREF(out_dtypes[i]);
1695
        /*
1696
         * Copy the dtype from 'op' if the type_num matches,
1697
         * to preserve metadata.
1698
         */
1699
        }
1700 1
        else if (op[i] != NULL &&
1701 1
                 PyArray_DESCR(op[i])->type_num == type_nums[i]) {
1702 1
            out_dtypes[i] = ensure_dtype_nbo(PyArray_DESCR(op[i]));
1703
        }
1704
        /*
1705
         * For outputs, copy the dtype from op[0] if the type_num
1706
         * matches, similarly to preserve metadata.
1707
         */
1708 1
        else if (i >= nin && op[0] != NULL &&
1709 1
                            PyArray_DESCR(op[0])->type_num == type_nums[i]) {
1710 1
            out_dtypes[i] = ensure_dtype_nbo(PyArray_DESCR(op[0]));
1711
        }
1712
        /* Otherwise create a plain descr from the type number */
1713
        else {
1714 1
            out_dtypes[i] = PyArray_DescrFromType(type_nums[i]);
1715
        }
1716

1717 1
        if (out_dtypes[i] == NULL) {
1718
            goto fail;
1719
        }
1720
    }
1721

1722
    return 0;
1723

1724 0
fail:
1725 0
    while (--i >= 0) {
1726 0
        Py_DECREF(out_dtypes[i]);
1727 0
        out_dtypes[i] = NULL;
1728
    }
1729
    return -1;
1730
}
1731

1732
/*
1733
 * Does a search through the arguments and the loops
1734
 */
1735
static int
1736 1
linear_search_userloop_type_resolver(PyUFuncObject *self,
1737
                        PyArrayObject **op,
1738
                        NPY_CASTING input_casting,
1739
                        NPY_CASTING output_casting,
1740
                        int any_object,
1741
                        int use_min_scalar,
1742
                        PyArray_Descr **out_dtype,
1743
                        int *out_no_castable_output,
1744
                        char *out_err_src_typecode,
1745
                        char *out_err_dst_typecode)
1746
{
1747 1
    npy_intp i, nop = self->nin + self->nout;
1748

1749
    /* Use this to try to avoid repeating the same userdef loop search */
1750 1
    int last_userdef = -1;
1751

1752 1
    for (i = 0; i < nop; ++i) {
1753
        int type_num;
1754

1755
        /* no more ufunc arguments to check */
1756 1
        if (op[i] == NULL) {
1757
            break;
1758
        }
1759

1760 1
        type_num = PyArray_DESCR(op[i])->type_num;
1761 1
        if (type_num != last_userdef &&
1762 1
                (PyTypeNum_ISUSERDEF(type_num) || type_num == NPY_VOID)) {
1763
            PyObject *key, *obj;
1764

1765 1
            last_userdef = type_num;
1766

1767 1
            key = PyLong_FromLong(type_num);
1768 1
            if (key == NULL) {
1769
                return -1;
1770
            }
1771 1
            obj = PyDict_GetItemWithError(self->userloops, key);
1772 1
            Py_DECREF(key);
1773 1
            if (obj == NULL && PyErr_Occurred()){
1774
                return -1;
1775
            }
1776 1
            else if (obj == NULL) {
1777 0
                continue;
1778
            }
1779 1
            PyUFunc_Loop1d *funcdata = PyCapsule_GetPointer(obj, NULL);
1780 1
            if (funcdata == NULL) {
1781
                return -1;
1782
            }
1783 0
            for (; funcdata != NULL; funcdata = funcdata->next) {
1784 1
                int *types = funcdata->arg_types;
1785 1
                switch (ufunc_loop_matches(self, op,
1786
                            input_casting, output_casting,
1787
                            any_object, use_min_scalar,
1788
                            types, funcdata->arg_dtypes,
1789
                            out_no_castable_output, out_err_src_typecode,
1790
                            out_err_dst_typecode)) {
1791
                    /* Error */
1792
                    case -1:
1793
                        return -1;
1794
                    /* Found a match */
1795 1
                    case 1:
1796 1
                        set_ufunc_loop_data_types(self, op, out_dtype, types, funcdata->arg_dtypes);
1797 1
                        return 1;
1798
                }
1799
            }
1800
        }
1801
    }
1802

1803
    /* Didn't find a match */
1804
    return 0;
1805
}
1806

1807
/*
1808
 * Does a search through the arguments and the loops
1809
 */
1810
static int
1811 1
type_tuple_userloop_type_resolver(PyUFuncObject *self,
1812
                        int n_specified,
1813
                        int *specified_types,
1814
                        PyArrayObject **op,
1815
                        NPY_CASTING casting,
1816
                        int any_object,
1817
                        int use_min_scalar,
1818
                        PyArray_Descr **out_dtype)
1819
{
1820 1
    int i, j, nin = self->nin, nop = nin + self->nout;
1821

1822
    /* Use this to try to avoid repeating the same userdef loop search */
1823 1
    int last_userdef = -1;
1824

1825 1
    int no_castable_output = 0;
1826 1
    char err_src_typecode = '-', err_dst_typecode = '-';
1827

1828 1
    for (i = 0; i < nin; ++i) {
1829 1
        int type_num = PyArray_DESCR(op[i])->type_num;
1830 1
        if (type_num != last_userdef && PyTypeNum_ISUSERDEF(type_num)) {
1831
            PyObject *key, *obj;
1832

1833 1
            last_userdef = type_num;
1834

1835 1
            key = PyLong_FromLong(type_num);
1836 1
            if (key == NULL) {
1837
                return -1;
1838
            }
1839 1
            obj = PyDict_GetItemWithError(self->userloops, key);
1840 1
            Py_DECREF(key);
1841 1
            if (obj == NULL && PyErr_Occurred()){
1842
                return -1;
1843
            }
1844 1
            else if (obj == NULL) {
1845 0
                continue;
1846
            }
1847

1848 1
            PyUFunc_Loop1d *funcdata = PyCapsule_GetPointer(obj, NULL);
1849 1
            if (funcdata == NULL) {
1850
                return -1;
1851
            }
1852 1
            for (; funcdata != NULL; funcdata = funcdata->next) {
1853 1
                int *types = funcdata->arg_types;
1854 1
                int matched = 1;
1855

1856 1
                if (n_specified == nop) {
1857 1
                    for (j = 0; j < nop; ++j) {
1858 1
                        if (types[j] != specified_types[j] &&
1859
                                    specified_types[j] != NPY_NOTYPE) {
1860
                            matched = 0;
1861
                            break;
1862
                        }
1863
                    }
1864
                } else {
1865 0
                    if (types[nin] != specified_types[0]) {
1866
                        matched = 0;
1867
                    }
1868
                }
1869 1
                if (!matched) {
1870 1
                    continue;
1871
                }
1872

1873 0
                switch (ufunc_loop_matches(self, op,
1874
                            casting, casting,
1875
                            any_object, use_min_scalar,
1876
                            types, NULL,
1877
                            &no_castable_output, &err_src_typecode,
1878
                            &err_dst_typecode)) {
1879
                    /* It works */
1880 0
                    case 1:
1881 0
                        set_ufunc_loop_data_types(self, op,
1882
                            out_dtype, types, NULL);
1883 0
                        return 1;
1884
                    /* Didn't match */
1885 0
                    case 0:
1886 0
                        PyErr_Format(PyExc_TypeError,
1887
                             "found a user loop for ufunc '%s' "
1888
                             "matching the type-tuple, "
1889
                             "but the inputs and/or outputs could not be "
1890
                             "cast according to the casting rule",
1891
                             ufunc_get_name_cstr(self));
1892 0
                        return -1;
1893
                    /* Error */
1894
                    case -1:
1895
                        return -1;
1896
                }
1897
            }
1898
        }
1899
    }
1900

1901
    /* Didn't find a match */
1902
    return 0;
1903
}
1904

1905

1906
/*
1907
 * Does a linear search for the best inner loop of the ufunc.
1908
 *
1909
 * Note that if an error is returned, the caller must free the non-zero
1910
 * references in out_dtype.  This function does not do its own clean-up.
1911
 */
1912
NPY_NO_EXPORT int
1913 1
linear_search_type_resolver(PyUFuncObject *self,
1914
                        PyArrayObject **op,
1915
                        NPY_CASTING input_casting,
1916
                        NPY_CASTING output_casting,
1917
                        int any_object,
1918
                        PyArray_Descr **out_dtype)
1919
{
1920 1
    npy_intp i, j, nin = self->nin, nop = nin + self->nout;
1921
    int types[NPY_MAXARGS];
1922
    const char *ufunc_name;
1923 1
    int no_castable_output = 0;
1924
    int use_min_scalar;
1925

1926
    /* For making a better error message on coercion error */
1927 1
    char err_dst_typecode = '-', err_src_typecode = '-';
1928

1929 1
    ufunc_name = ufunc_get_name_cstr(self);
1930

1931 1
    use_min_scalar = should_use_min_scalar(nin, op, 0, NULL);
1932

1933
    /* If the ufunc has userloops, search for them. */
1934 1
    if (self->userloops) {
1935 1
        switch (linear_search_userloop_type_resolver(self, op,
1936
                                input_casting, output_casting,
1937
                                any_object, use_min_scalar, out_dtype,
1938
                                &no_castable_output, &err_src_typecode,
1939
                                &err_dst_typecode)) {
1940
            /* Error */
1941
            case -1:
1942
                return -1;
1943
            /* A loop was found */
1944 1
            case 1:
1945 1
                return 0;
1946
        }
1947
    }
1948

1949
    /*
1950
     * Determine the UFunc loop.  This could in general be *much* faster,
1951
     * and a better way to implement it might be for the ufunc to
1952
     * provide a function which gives back the result type and inner
1953
     * loop function.
1954
     *
1955
     * A default fast mechanism could be provided for functions which
1956
     * follow the most typical pattern, when all functions have signatures
1957
     * "xx...x -> x" for some built-in data type x, as follows.
1958
     *  - Use PyArray_ResultType to get the output type
1959
     *  - Look up the inner loop in a table based on the output type_num
1960
     *
1961
     * The method for finding the loop in the previous code did not
1962
     * appear consistent (as noted by some asymmetry in the generated
1963
     * coercion tables for np.add).
1964
     */
1965 1
    no_castable_output = 0;
1966 1
    for (i = 0; i < self->ntypes; ++i) {
1967 1
        char *orig_types = self->types + i*self->nargs;
1968

1969
        /* Copy the types into an int array for matching */
1970 1
        for (j = 0; j < nop; ++j) {
1971 1
            types[j] = orig_types[j];
1972
        }
1973

1974 1
        switch (ufunc_loop_matches(self, op,
1975
                    input_casting, output_casting,
1976
                    any_object, use_min_scalar,
1977
                    types, NULL,
1978
                    &no_castable_output, &err_src_typecode,
1979
                    &err_dst_typecode)) {
1980
            /* Error */
1981
            case -1:
1982
                return -1;
1983
            /* Found a match */
1984 1
            case 1:
1985 1
                set_ufunc_loop_data_types(self, op, out_dtype, types, NULL);
1986 1
                return 0;
1987
        }
1988
    }
1989

1990
    /* If no function was found, throw an error */
1991 1
    if (no_castable_output) {
1992 1
        PyErr_Format(PyExc_TypeError,
1993
                "ufunc '%s' output (typecode '%c') could not be coerced to "
1994
                "provided output parameter (typecode '%c') according "
1995
                "to the casting rule '%s'",
1996
                ufunc_name, err_src_typecode, err_dst_typecode,
1997
                npy_casting_to_string(output_casting));
1998
    }
1999
    else {
2000
        /*
2001
         * TODO: We should try again if the casting rule is same_kind
2002
         *       or unsafe, and look for a function more liberally.
2003
         */
2004 1
        PyErr_Format(PyExc_TypeError,
2005
                "ufunc '%s' not supported for the input types, and the "
2006
                "inputs could not be safely coerced to any supported "
2007
                "types according to the casting rule '%s'",
2008
                ufunc_name,
2009
                npy_casting_to_string(input_casting));
2010
    }
2011

2012
    return -1;
2013
}
2014

2015
/*
2016
 * Does a linear search for the inner loop of the ufunc specified by type_tup.
2017
 *
2018
 * Note that if an error is returned, the caller must free the non-zero
2019
 * references in out_dtype.  This function does not do its own clean-up.
2020
 */
2021
NPY_NO_EXPORT int
2022 1
type_tuple_type_resolver(PyUFuncObject *self,
2023
                        PyObject *type_tup,
2024
                        PyArrayObject **op,
2025
                        NPY_CASTING casting,
2026
                        int any_object,
2027
                        PyArray_Descr **out_dtype)
2028
{
2029 1
    npy_intp i, j, n, nin = self->nin, nop = nin + self->nout;
2030 1
    int n_specified = 0;
2031
    int specified_types[NPY_MAXARGS], types[NPY_MAXARGS];
2032
    const char *ufunc_name;
2033
    int no_castable_output, use_min_scalar;
2034

2035
    /* For making a better error message on coercion error */
2036 1
    char err_dst_typecode = '-', err_src_typecode = '-';
2037

2038 1
    ufunc_name = ufunc_get_name_cstr(self);
2039

2040 1
    use_min_scalar = should_use_min_scalar(nin, op, 0, NULL);
2041

2042
    /* Fill in specified_types from the tuple or string */
2043 1
    if (PyTuple_Check(type_tup)) {
2044 1
        int nonecount = 0;
2045 1
        n = PyTuple_GET_SIZE(type_tup);
2046 1
        if (n != 1 && n != nop) {
2047 0
            PyErr_Format(PyExc_ValueError,
2048
                         "a type-tuple must be specified "
2049
                         "of length 1 or %d for ufunc '%s'", (int)nop,
2050
                         ufunc_get_name_cstr(self));
2051 0
            return -1;
2052
        }
2053

2054 1
        for (i = 0; i < n; ++i) {
2055 1
            PyObject *item = PyTuple_GET_ITEM(type_tup, i);
2056 1
            if (item == Py_None) {
2057 1
                specified_types[i] = NPY_NOTYPE;
2058 1
                ++nonecount;
2059
            }
2060
            else {
2061 1
                PyArray_Descr *dtype = NULL;
2062 1
                if (!PyArray_DescrConverter(item, &dtype)) {
2063 0
                    return -1;
2064
                }
2065 1
                specified_types[i] = dtype->type_num;
2066 1
                Py_DECREF(dtype);
2067
            }
2068
        }
2069

2070 1
        if (nonecount == n) {
2071 0
            PyErr_SetString(PyExc_ValueError,
2072
                    "the type-tuple provided to the ufunc "
2073
                    "must specify at least one none-None dtype");
2074 0
            return -1;
2075
        }
2076

2077 1
        n_specified = n;
2078
    }
2079 1
    else if (PyBytes_Check(type_tup) || PyUnicode_Check(type_tup)) {
2080
        Py_ssize_t length;
2081
        char *str;
2082 1
        PyObject *str_obj = NULL;
2083

2084 1
        if (PyUnicode_Check(type_tup)) {
2085 1
            str_obj = PyUnicode_AsASCIIString(type_tup);
2086 1
            if (str_obj == NULL) {
2087 0
                return -1;
2088
            }
2089
            type_tup = str_obj;
2090
        }
2091

2092 1
        if (PyBytes_AsStringAndSize(type_tup, &str, &length) < 0) {
2093 0
            Py_XDECREF(str_obj);
2094
            return -1;
2095
        }
2096 1
        if (length != 1 && (length != nop + 2 ||
2097 1
                                str[nin] != '-' || str[nin+1] != '>')) {
2098 0
            PyErr_Format(PyExc_ValueError,
2099
                                 "a type-string for %s, "   \
2100
                                 "requires 1 typecode, or "
2101
                                 "%d typecode(s) before " \
2102
                                 "and %d after the -> sign",
2103
                                 ufunc_get_name_cstr(self),
2104
                                 self->nin, self->nout);
2105 0
            Py_XDECREF(str_obj);
2106
            return -1;
2107
        }
2108 1
        if (length == 1) {
2109
            PyArray_Descr *dtype;
2110 1
            n_specified = 1;
2111 1
            dtype = PyArray_DescrFromType(str[0]);
2112 1
            if (dtype == NULL) {
2113 0
                Py_XDECREF(str_obj);
2114
                return -1;
2115
            }
2116 1
            specified_types[0] = dtype->type_num;
2117 1
            Py_DECREF(dtype);
2118
        }
2119
        else {
2120
            PyArray_Descr *dtype;
2121 1
            n_specified = (int)nop;
2122

2123 1
            for (i = 0; i < nop; ++i) {
2124 1
                npy_intp istr = i < nin ? i : i+2;
2125

2126 1
                dtype = PyArray_DescrFromType(str[istr]);
2127 1
                if (dtype == NULL) {
2128 0
                    Py_XDECREF(str_obj);
2129
                    return -1;
2130
                }
2131 1
                specified_types[i] = dtype->type_num;
2132 1
                Py_DECREF(dtype);
2133
            }
2134
        }
2135 1
        Py_XDECREF(str_obj);
2136
    }
2137

2138
    /* If the ufunc has userloops, search for them. */
2139 1
    if (self->userloops) {
2140 1
        switch (type_tuple_userloop_type_resolver(self,
2141
                        n_specified, specified_types,
2142
                        op, casting,
2143
                        any_object, use_min_scalar,
2144
                        out_dtype)) {
2145
            /* Error */
2146
            case -1:
2147
                return -1;
2148
            /* Found matching loop */
2149 0
            case 1:
2150 0
                return 0;
2151
        }
2152
    }
2153

2154 1
    for (i = 0; i < self->ntypes; ++i) {
2155 1
        char *orig_types = self->types + i*self->nargs;
2156

2157
        /* Copy the types into an int array for matching */
2158 1
        for (j = 0; j < nop; ++j) {
2159 1
            types[j] = orig_types[j];
2160
        }
2161

2162 1
        if (n_specified == nop) {
2163 1
            for (j = 0; j < nop; ++j) {
2164 1
                if (types[j] != specified_types[j] &&
2165
                        specified_types[j] != NPY_NOTYPE) {
2166
                    break;
2167
                }
2168
            }
2169 1
            if (j < nop) {
2170
                /* no match */
2171 1
                continue;
2172
            }
2173
        }
2174 1
        else if (types[nin] != specified_types[0]) {
2175
            /* no match */
2176 1
            continue;
2177
        }
2178

2179 1
        switch (ufunc_loop_matches(self, op,
2180
                    casting, casting,
2181
                    any_object, use_min_scalar,
2182
                    types, NULL,
2183
                    &no_castable_output, &err_src_typecode,
2184
                    &err_dst_typecode)) {
2185
            case -1:
2186
                /* Error */
2187
                return -1;
2188 1
            case 0:
2189
                /* Cannot cast inputs */
2190 1
                continue;
2191 1
            case 1:
2192
                /* Success */
2193 1
                set_ufunc_loop_data_types(self, op, out_dtype, types, NULL);
2194 1
                return 0;
2195
        }
2196
    }
2197

2198
    /* If no function was found, throw an error */
2199 1
    PyErr_Format(PyExc_TypeError,
2200
            "No loop matching the specified signature and casting "
2201
            "was found for ufunc %s", ufunc_name);
2202

2203 1
    return -1;
2204
}
2205

2206
NPY_NO_EXPORT int
2207 1
PyUFunc_DivmodTypeResolver(PyUFuncObject *ufunc,
2208
                                NPY_CASTING casting,
2209
                                PyArrayObject **operands,
2210
                                PyObject *type_tup,
2211
                                PyArray_Descr **out_dtypes)
2212
{
2213
    int type_num1, type_num2;
2214
    int i;
2215

2216 1
    type_num1 = PyArray_DESCR(operands[0])->type_num;
2217 1
    type_num2 = PyArray_DESCR(operands[1])->type_num;
2218

2219
    /* Use the default when datetime and timedelta are not involved */
2220 1
    if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
2221 1
        return PyUFunc_DefaultTypeResolver(ufunc, casting, operands,
2222
                    type_tup, out_dtypes);
2223
    }
2224 1
    if (type_num1 == NPY_TIMEDELTA) {
2225 1
        if (type_num2 == NPY_TIMEDELTA) {
2226 1
            out_dtypes[0] = PyArray_PromoteTypes(PyArray_DESCR(operands[0]),
2227
                                                PyArray_DESCR(operands[1]));
2228 1
            out_dtypes[1] = out_dtypes[0];
2229 1
            Py_INCREF(out_dtypes[1]);
2230 1
            out_dtypes[2] = PyArray_DescrFromType(NPY_LONGLONG);
2231 1
            out_dtypes[3] = out_dtypes[0];
2232 1
            Py_INCREF(out_dtypes[3]);
2233
        }
2234
        else {
2235 0
            return raise_binary_type_reso_error(ufunc, operands);
2236
        }
2237
    }
2238
    else {
2239 0
        return raise_binary_type_reso_error(ufunc, operands);
2240
    }
2241

2242
    /* Check against the casting rules */
2243 1
    if (PyUFunc_ValidateCasting(ufunc, casting, operands, out_dtypes) < 0) {
2244 0
        for (i = 0; i < 4; ++i) {
2245 0
            Py_DECREF(out_dtypes[i]);
2246 0
            out_dtypes[i] = NULL;
2247
        }
2248
        return -1;
2249
    }
2250

2251
    return 0;
2252
}

Read our documentation on viewing source code .

Loading