MAINT: check enum at API call, not in provate function

mattip · mattip · commit 26079cf123f4 · 2019-09-20T17:57:54.000+03:00
diff --git a/numpy/core/src/common/npy_partition.h.src b/numpy/core/src/common/npy_partition.h.src
@@ -113,9 +113,6 @@ get_argpartition_func(int type, NPY_SELECTKIND which)
     npy_intp i;
     npy_intp ntypes = ARRAY_SIZE(_part_map);
 
-    if (which >= NPY_NSELECTS) {
-        return NULL;
-    }
     for (i = 0; i < ntypes; i++) {
         if (type == _part_map[i].typenum) {
             return _part_map[i].argpart[which];
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
@@ -1336,7 +1336,11 @@ PyArray_ArgPartition(PyArrayObject *op, PyArrayObject *ktharray, int axis,
     PyArray_ArgSortFunc *argsort;
     PyObject *ret;
 
-    if (which < 0 || which >= NPY_NSELECTS) {
+    /*
+     * As a C-exported function, enum NPY_SELECTKIND loses its enum property
+     * Check the values to make sure they are in range
+     */
+    if ((int)which < 0 || (int)which >= NPY_NSELECTS) {
         PyErr_SetString(PyExc_ValueError,
                         "not a valid partition kind");
         return NULL;
