@@ -70,7 +70,37 @@ promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
7070static int
7171add_ufunc_loop (PyUFuncObject * ufunc , PyObject * info , int ignore_duplicate )
7272{
73- assert (PyTuple_CheckExact (info ) && PyTuple_GET_SIZE (info ) == 2 );
73+ /*
74+ * Validate the info object, this should likely move to to a different
75+ * entry-point in the future (and is mostly unnecessary currently).
76+ */
77+ if (!PyTuple_CheckExact (info ) || PyTuple_GET_SIZE (info ) != 2 ) {
78+ PyErr_SetString (PyExc_TypeError ,
79+ "Info must be a tuple: "
80+ "(tuple of DTypes or None, ArrayMethod or promoter)" );
81+ return -1 ;
82+ }
83+ PyObject * DType_tuple = PyTuple_GetItem (info , 0 );
84+ if (PyTuple_GET_SIZE (DType_tuple ) != ufunc -> nargs ) {
85+ PyErr_SetString (PyExc_TypeError ,
86+ "DType tuple length does not match ufunc number of operands" );
87+ return -1 ;
88+ }
89+ for (Py_ssize_t i = 0 ; i < PyTuple_GET_SIZE (DType_tuple ); i ++ ) {
90+ PyObject * item = PyTuple_GET_ITEM (DType_tuple , i );
91+ if (item != Py_None
92+ && !PyObject_TypeCheck (item , & PyArrayDTypeMeta_Type )) {
93+ PyErr_SetString (PyExc_TypeError ,
94+ "DType tuple may only contain None and DType classes" );
95+ return -1 ;
96+ }
97+ }
98+ if (!PyObject_TypeCheck (PyTuple_GET_ITEM (info , 1 ), & PyArrayMethod_Type )) {
99+ /* Must also accept promoters in the future. */
100+ PyErr_SetString (PyExc_TypeError ,
101+ "Second argument to info must be an ArrayMethod or promoter" );
102+ return -1 ;
103+ }
74104
75105 if (ufunc -> _loops == NULL ) {
76106 ufunc -> _loops = PyList_New (0 );
@@ -79,8 +109,6 @@ add_ufunc_loop(PyUFuncObject *ufunc, PyObject *info, int ignore_duplicate)
79109 }
80110 }
81111
82- PyObject * DType_tuple = PyTuple_GetItem (info , 0 );
83-
84112 PyObject * loops = ufunc -> _loops ;
85113 Py_ssize_t length = PyList_Size (loops );
86114 for (Py_ssize_t i = 0 ; i < length ; i ++ ) {
0 commit comments