8000 Merge pull request #10303 from mhvk/array-wrap-prepare-cleanup · eric-wieser/numpy@7f74525 · GitHub
[go: up one dir, main page]

Skip to content

Commit 7f74525

Browse files
authored
Merge pull request numpy#10303 from mhvk/array-wrap-prepare-cleanup
MAINT: Array wrap/prepare identification cleanup
2 parents f1d81d8 + c9179a9 commit 7f74525

File tree

1 file changed

+72
-112
lines changed

1 file changed

+72
-112
lines changed

numpy/core/src/umath/ufunc_object.c

Lines changed: 72 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -126,94 +126,94 @@ PyUFunc_clearfperr()
126126
}
127127

128128
/*
129-
* This function analyzes the input arguments
130-
* and determines an appropriate __array_prepare__ function to call
131-
* for the outputs.
132-
* Assumes subok is already true if check_subok is false.
133-
*
134-
* If an output argument is provided, then it is prepped
135-
* with its own __array_prepare__ not with the one determined by
136-
* the input arguments.
137-
*
138-
* if the provided output argument is already an ndarray,
139-
* the prepping function is None (which means no prepping will
140-
* be done --- not even PyArray_Return).
141-
*
142-
* A NULL is placed in output_prep for outputs that
143-
* should just have PyArray_Return called.
129+
* This function analyzes the input arguments and determines an appropriate
130+
* method (__array_prepare__ or __array_wrap__) function to call, taking it
131+
* from the input with the highest priority. Return NULL if no argument
132+
* defines the method.
144133
*/
145-
static void
146-
_find_array_prepare(PyObject *args, PyObject *kwds,
147-
PyObject **output_prep, int nin, int nout,
148-
int check_subok)
134+
static PyObject*
135+
_find_array_method(PyObject *args, int nin, PyObject *method_name)
149136
{
150-
Py_ssize_t nargs;
151-
int i;
152-
int np = 0;
153-
PyObject *with_prep[NPY_MAXARGS], *preps[NPY_MAXARGS];
154-
PyObject *obj, *prep = NULL;
155-
156-
/*
157-
* If a 'subok' parameter is passed and isn't True, don't wrap
158-
* if check_subok is false it assumed subok in kwds keyword is True
159-
*/
160-
if (check_subok && kwds != NULL &&
161-
(obj = PyDict_GetItem(kwds, npy_um_str_subok)) != NULL) {
162-
if (obj != Py_True) {
163-
for (i = 0; i < nout; i++) {
164-
output_prep[i] = NULL;
165-
}
166-
return;
167-
}
168-
}
137+
int i, n_methods;
138+
PyObject *obj;
139+
PyObject *with_method[NPY_MAXARGS], *methods[NPY_MAXARGS];
140+
PyObject *method = NULL;
169141

170-
nargs = PyTuple_GET_SIZE(args);
142+
n_methods = 0;
171143
for (i = 0; i < nin; i++) {
172144
obj = PyTuple_GET_ITEM(args, i);
173145
if (PyArray_CheckExact(obj) || PyArray_IsAnyScalar(obj)) {
174146
continue;
175147
}
176-
prep = PyObject_GetAttr(obj, npy_um_str_array_prepare);
177-
if (prep) {
178-
if (PyCallable_Check(prep)) {
179-
with_prep[np] = obj;
180-
preps[np] = prep;
181-
++np;
148+
method = PyObject_GetAttr(obj, method_name);
149+
if (method) {
150+
if (PyCallable_Check(method)) {
151+
with_method[n_methods] = obj;
152+
methods[n_methods] = method;
153+
++n_methods;
182154
}
183155
else {
184-
Py_DECREF(prep);
185-
prep = NULL;
156+
Py_DECREF(method);
157+
method = NULL;
186158
}
187159
}
188160
else {
189161
PyErr_Clear();
190162
}
191163
}
192-
if (np > 0) {
193-
/* If we have some preps defined, find the one of highest priority */
194-
prep = preps[0];
195-
if (np > 1) {
196-
double maxpriority = PyArray_GetPriority(with_prep[0],
197-
NPY_PRIORITY);
198-
for (i = 1; i < np; ++i) {
199-
double priority = PyArray_GetPriority(with_prep[i],
200-
NPY_PRIORITY);
164+
if (n_methods > 0) {
165+
/* If we have some methods defined, find the one of highest priority */
166+
method = methods[0];
167+
if (n_methods > 1) {
168+
double maxpriority = PyArray_GetPriority(with_method[0],
169+
NPY_PRIORITY);
170+
for (i = 1; i < n_methods; ++i) {
171+
double priority = PyArray_GetPriority(with_method[i],
172+
NPY_PRIORITY);
201173
if (priority > maxpriority) {
202174
maxpriority = priority;
203-
Py_DECREF(prep);
204-
prep = preps[i];
175+
Py_DECREF(method);
176+
method = methods[i];
205177
}
206178
else {
207-
Py_DECREF(preps[i]);
179+
Py_DECREF(methods[i]);
208180
}
209181
}
210182
}
211183
}
184+
return method;
185+
}
186+
187+
/*
188+
* This function analyzes the input arguments
189+
* and determines an appropriate __array_prepare__ function to call
190+
* for the outputs.
191+
*
192+
* If an output argument is provided, then it is prepped
193+
* with its own __array_prepare__ not with the one determined by
194+
* the input arguments.
195+
*
196+
* if the provided output argument is already an ndarray,
197+
* the prepping function is None (which means no prepping will
198+
* be done --- not even PyArray_Return).
199+
*
200+
* A NULL is placed in output_prep for outputs that
201+
* should just have PyArray_Return called.
202+
*/
203+
static void
204+
_find_array_prepare(PyObject *args, PyObject *kwds,
205+
PyObject **output_prep, int nin, int nout)
206+
{
207+
Py_ssize_t nargs;
208+
int i;
209+
PyObject *obj, *prep;
212210

213211
/*
214-
* Here prep is the prepping function determined from the
215-
* input arrays (could be NULL).
216-
*
212+
* Determine the prepping function given by the input arrays
213+
* (could be NULL).
214+
*/
215+
prep = _find_array_method(args, nin, npy_um_str_array_prepare);
216+
/*
217217
* For all the output arrays decide what to do.
218218
*
219219
* 1) Use the prep function determined from the input arrays
@@ -225,6 +225,7 @@ _find_array_prepare(PyObject *args, PyObject *kwds,
225225
* exact ndarray so that no PyArray_Return is
226226
* done in that case.
227227
*/
228+
nargs = PyTuple_GET_SIZE(args);
228229
for (i = 0; i < nout; i++) {
229230
int j = nin + i;
230231
int incref = 1;
@@ -2101,7 +2102,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
21012102
* Get the appropriate __array_prepare__ function to call
21022103
* for each output
21032104
*/
2104-
_find_array_prepare(args, kwds, arr_prep, nin, nout, 0);
2105+
_find_array_prepare(args, kwds, arr_prep, nin, nout);
21052106

21062107
10000 /* Set up arr_prep_args if a prep function was needed */
21072108
for (i = 0; i < nout; ++i) {
@@ -2486,7 +2487,7 @@ PyUFunc_GenericFunction(PyUFuncObject *ufunc,
24862487
* Get the appropriate __array_prepare__ function to call
24872488
* for each output
24882489
*/
2489-
_find_array_prepare(args, kwds, arr_prep, nin, nout, 0);
2490+
_find_array_prepare(args, kwds, arr_prep, nin, nout);
24902491

24912492
/* Set up arr_prep_args if a prep function was needed */
24922493
for (i = 0; i < nout; ++i) {
@@ -3946,9 +3947,8 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
39463947
{
39473948
Py_ssize_t nargs;
39483949
int i, idx_offset, start_idx;
3949-
int np = 0;
3950-
PyObject *with_wrap[NPY_MAXARGS], *wraps[NPY_MAXARGS];
3951-
PyObject *obj, *wrap = NULL;
3950+
PyObject *obj;
3951+
PyObject *wrap = NULL;
39523952

39533953
/*
39543954
* If a 'subok' parameter is passed and isn't True, don't wrap but put None
@@ -3962,53 +3962,13 @@ _find_array_wrap(PyObject *args, PyObject *kwds,
39623962
}
39633963
}
39643964

3965-
3966-
for (i = 0; i < nin; i++) {
3967-
obj = PyTuple_GET_ITEM(args, i);
3968-
if (PyArray_CheckExact(obj) || PyArray_IsAnyScalar(obj)) {
3969-
continue;
3970-
}
3971-
wrap = PyObject_GetAttr(obj, npy_um_str_array_wrap);
3972-
if (wrap) {
3973-
if (PyCallable_Check(wrap)) {
3974-
with_wrap[np] = obj;
3975-
wraps[np] = wrap;
3976-
++np;
3977-
}
3978-
else {
3979-
Py_DECREF(wrap);
3980-
wrap = NULL;
3981-
}
3982-
}
3983-
else {
3984-
PyErr_Clear();
3985-
}
3986-
}
3987-
if (np > 0) {
3988-
/* If we have some wraps defined, find the one of highest priority */
3989-
wrap = wraps[0];
3990-
if (np > 1) {
3991-
double maxpriority = PyArray_GetPriority(with_wrap[0],
3992-
NPY_PRIORITY);
3993-
for (i = 1; i < np; ++i) {
3994-
double priority = PyArray_GetPriority(with_wrap[i],
3995-
NPY_PRIORITY);
3996-
if (priority > maxpriority) {
3997-
maxpriority = priority;
3998-
Py_DECREF(wrap);
3999-
wrap = wraps[i];
4000-
}
4001-
else {
4002-
Py_DECREF(wraps[i]);
4003-
}
4004-
}
4005-
}
4006-
}
3965+
/*
3966+
* Determine the wrapping function given by the input arrays
3967+
* (could be NULL).
3968+
*/
3969+
wrap = _find_array_method(args, nin, npy_um_str_array_wrap);
40073970

40083971
/*
4009-
* Here wrap is the wrapping function determined from the
4010-
* input arrays (could be NULL).
4011-
*
40123972
* For all the output arrays decide what to do.
40133973
*
40143974
* 1) Use the wrap function determined from the input arrays

0 commit comments

Comments
 (0)
0