Function PyUFunc_AddWrappingLoop

numpy/core/src/umath/wrapping_array_method.c:225–337

* Allows creating of a fairly lightweight wrapper around an existing ufunc * loop. The idea is mainly for units, as this is currently slightly limited * in that it enforces that you cannot use a loop from another ufunc. * * @param ufunc_obj * @param new_dtypes * @param wrapped_dtypes * @param translate_given_descrs See typedef comment * @param translate_loop_descrs See typedef comment *

 * @return 0 on success -1 on failure
 */
NPY_NO_EXPORT int
PyUFunc_AddWrappingLoop(PyObject *ufunc_obj,
        PyArray_DTypeMeta *new_dtypes[], PyArray_DTypeMeta *wrapped_dtypes[],
        translate_given_descrs_func *translate_given_descrs,
        translate_loop_descrs_func *translate_loop_descrs)
{
    int res = -1;
    PyUFuncObject *ufunc = (PyUFuncObject *)ufunc_obj;
    PyObject *wrapped_dt_tuple = NULL;
    PyObject *new_dt_tuple = NULL;
    PyArrayMethodObject *meth = NULL;

    if (!PyObject_TypeCheck(ufunc_obj, &PyUFunc_Type)) {
        PyErr_SetString(PyExc_TypeError,
                "ufunc object passed is not a ufunc!");
        return -1;
    }

    wrapped_dt_tuple = PyArray_TupleFromItems(
            ufunc->nargs, (PyObject **)wrapped_dtypes, 1);
    if (wrapped_dt_tuple == NULL) {
        goto finish;
    }

    PyArrayMethodObject *wrapped_meth = NULL;
    PyObject *loops = ufunc->_loops;
    Py_ssize_t length = PyList_Size(loops);
    for (Py_ssize_t i = 0; i < length; i++) {
        PyObject *item = PyList_GetItem(loops, i);
        PyObject *cur_DType_tuple = PyTuple_GetItem(item, 0);
        int cmp = PyObject_RichCompareBool(cur_DType_tuple, wrapped_dt_tuple, Py_EQ);
        if (cmp < 0) {
            goto finish;
        }
        if (cmp == 0) {
            continue;
        }
        wrapped_meth = (PyArrayMethodObject *)PyTuple_GET_ITEM(item, 1);
        if (!PyObject_TypeCheck(wrapped_meth, &PyArrayMethod_Type)) {
            PyErr_SetString(PyExc_TypeError,
                    "Matching loop was not an ArrayMethod.");
            goto finish;
        }
        break;
    }
    if (wrapped_meth == NULL) {
        PyErr_Format(PyExc_TypeError,
                "Did not find the to-be-wrapped loop in the ufunc with given "
                "DTypes. Received wrapping types: %S", wrapped_dt_tuple);
        goto finish;
    }

    PyType_Slot slots[] = {
        {NPY_METH_resolve_descriptors, &wrapping_method_resolve_descriptors},
        {_NPY_METH_get_loop, &wrapping_method_get_loop},
        {NPY_METH_get_reduction_initial,
            &wrapping_method_get_identity_function},
        {0, NULL}