Function PyArray_Where

numpy/core/src/multiarray/multiarraymodule.c:3368–3550

NUMPY_API * Where */

 * Where
 */
NPY_NO_EXPORT PyObject *
PyArray_Where(PyObject *condition, PyObject *x, PyObject *y)
{
    PyArrayObject *arr, *ax, *ay;
    PyObject *ret = NULL;

    arr = (PyArrayObject *)PyArray_FROM_O(condition);
    if (arr == NULL) {
        return NULL;
    }
    if ((x == NULL) && (y == NULL)) {
        ret = PyArray_Nonzero(arr);
        Py_DECREF(arr);
        return ret;
    }
    if ((x == NULL) || (y == NULL)) {
        Py_DECREF(arr);
        PyErr_SetString(PyExc_ValueError,
                "either both or neither of x and y should be given");
        return NULL;
    }

    NPY_cast_info cast_info = {.func = NULL};

    ax = (PyArrayObject*)PyArray_FROM_O(x);
    ay = (PyArrayObject*)PyArray_FROM_O(y);
    if (ax == NULL || ay == NULL) {
        goto fail;
    }
    else {
        npy_uint32 flags = NPY_ITER_EXTERNAL_LOOP | NPY_ITER_BUFFERED |
                           NPY_ITER_REFS_OK | NPY_ITER_ZEROSIZE_OK;
        PyArrayObject * op_in[4] = {
            NULL, arr, ax, ay
        };
        npy_uint32 op_flags[4] = {
            NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE | NPY_ITER_NO_SUBTYPE,
            NPY_ITER_READONLY,
            NPY_ITER_READONLY | NPY_ITER_ALIGNED,
            NPY_ITER_READONLY | NPY_ITER_ALIGNED
        };
        PyArray_Descr * common_dt = PyArray_ResultType(2, &op_in[0] + 2,
                                                       0, NULL);
        PyArray_Descr * op_dt[4] = {common_dt, PyArray_DescrFromType(NPY_BOOL),
                                    common_dt, common_dt};
        NpyIter * iter;
        NPY_BEGIN_THREADS_DEF;

        if (common_dt == NULL || op_dt[1] == NULL) {
            Py_XDECREF(op_dt[1]);
            Py_XDECREF(common_dt);
            goto fail;
        }
        iter =  NpyIter_MultiNew(4, op_in, flags,
                                 NPY_KEEPORDER, NPY_UNSAFE_CASTING,
                                 op_flags, op_dt);
        Py_DECREF(op_dt[1]);
        Py_DECREF(common_dt);