BUG: Prevent hang when sorting dtype=S0 arrays

3 files changed, 87 insertions, 4 deletions

diff --git a/numpy/core/src/npysort/mergesort.c.src b/numpy/core/src/npysort/mergesort.c.src
index fc82e2135..6f659617a 100644
--- a/numpy/core/src/npysort/mergesort.c.src
+++ b/numpy/core/src/npysort/mergesort.c.src
@@ -254,6 +254,11 @@ mergesort_@suff@(void *start, npy_intp num, void *varr)
     @type@ *pl, *pr, *pw, *vp;
     int err = 0;
 
+    /* Items that have zero size don't make sense to sort */
+    if (elsize == 0) {
+        return 0;
+    }
+
     pl = start;
     pr = pl + num*len;
     pw = malloc((num/2) * elsize);
@@ -329,6 +334,11 @@ amergesort_@suff@(void *v, npy_intp *tosort, npy_intp num, void *varr)
     size_t len = elsize / sizeof(@type@);
     npy_intp *pl, *pr, *pw;
 
+    /* Items that have zero size don't make sense to sort */
+    if (elsize == 0) {
+        return 0;
+    }
+
     pl = tosort;
     pr = pl + num;
     pw = malloc((num/2) * sizeof(npy_intp));
@@ -405,10 +415,18 @@ npy_mergesort(void *start, npy_intp num, void *varr)
     PyArray_CompareFunc *cmp = PyArray_DESCR(arr)->f->compare;
     char *pl = start;
     char *pr = pl + num*elsize;
-    char *pw = malloc((num >> 1) *elsize);
-    char *vp = malloc(elsize);
+    char *pw;
+    char *vp;
     int err = -NPY_ENOMEM;
 
+    /* Items that have zero size don't make sense to sort */
+    if (elsize == 0) {
+        return 0;
+    }
+
+    pw = malloc((num >> 1) *elsize);
+    vp = malloc(elsize);
+
     if (pw != NULL && vp != NULL) {
         npy_mergesort0(pl, pr, pw, vp, elsize, cmp, arr);
         err = 0;
@@ -475,6 +493,11 @@ npy_amergesort(void *v, npy_intp *tosort, npy_intp num, void *varr)
     PyArray_CompareFunc *cmp = PyArray_DESCR(arr)->f->compare;
     npy_intp *pl, *pr, *pw;
 
+    /* Items that have zero size don't make sense to sort */
+    if (elsize == 0) {
+        return 0;
+    }
+
     pl = tosort;
     pr = pl + num;
     pw = malloc((num >> 1) * sizeof(npy_intp));
diff --git a/numpy/core/src/npysort/quicksort.c.src b/numpy/core/src/npysort/quicksort.c.src
index 2b6e2ed1c..ff0e8a149 100644
--- a/numpy/core/src/npysort/quicksort.c.src
+++ b/numpy/core/src/npysort/quicksort.c.src
@@ -258,7 +258,7 @@ quicksort_@suff@(void *start, npy_intp num, void *varr)
 {
     PyArrayObject *arr = varr;
     const size_t len = PyArray_ITEMSIZE(arr)/sizeof(@type@);
-    @type@ *vp = malloc(PyArray_ITEMSIZE(arr));
+    @type@ *vp;
     @type@ *pl = start;
     @type@ *pr = pl + (num - 1)*len;
     @type@ *stack[PYA_QS_STACK], **sptr = stack, *pm, *pi, *pj, *pk;
@@ -266,6 +266,12 @@ quicksort_@suff@(void *start, npy_intp num, void *varr)
     int * psdepth = depth;
     int cdepth = npy_get_msb(num) * 2;
 
+    /* Items that have zero size don't make sense to sort */
+    if (len == 0) {
+        return 0;
+    }
+
+    vp = malloc(PyArray_ITEMSIZE(arr));
     if (vp == NULL) {
         return -NPY_ENOMEM;
     }
@@ -351,6 +357,11 @@ aquicksort_@suff@(void *vv, npy_intp* tosort, npy_intp num, void *varr)
     int * psdepth = depth;
     int cdepth = npy_get_msb(num) * 2;
 
+    /* Items that have zero size don't make sense to sort */
+    if (len == 0) {
+        return 0;
+    }
+
     for (;;) {
         if (NPY_UNLIKELY(cdepth < 0)) {
             aheapsort_@suff@(vv, pl, pr - pl + 1, varr);
@@ -429,7 +440,7 @@ npy_quicksort(void *start, npy_intp num, void *varr)
     PyArrayObject *arr = varr;
     npy_intp elsize = PyArray_ITEMSIZE(arr);
     PyArray_CompareFunc *cmp = PyArray_DESCR(arr)->f->compare;
-    char *vp = malloc(elsize);
+    char *vp;
     char *pl = start;
     char *pr = pl + (num - 1)*elsize;
     char *stack[PYA_QS_STACK];
@@ -439,6 +450,12 @@ npy_quicksort(void *start, npy_intp num, void *varr)
     int * psdepth = depth;
     int cdepth = npy_get_msb(num) * 2;
 
+    /* Items that have zero size don't make sense to sort */
+    if (elsize == 0) {
+        return 0;
+    }
+
+    vp = malloc(elsize);
     if (vp == NULL) {
         return -NPY_ENOMEM;
     }
@@ -539,6 +556,11 @@ npy_aquicksort(void *vv, npy_intp* tosort, npy_intp num, void *varr)
     int * psdepth = depth;
     int cdepth = npy_get_msb(num) * 2;
 
+    /* Items that have zero size don't make sense to sort */
+    if (elsize == 0) {
+        return 0;
+    }
+
     for (;;) {
         if (NPY_UNLIKELY(cdepth < 0)) {
             npy_aheapsort(vv, pl, pr - pl + 1, varr);
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index c28a72150..9703fdbb5 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -1204,6 +1204,44 @@ class TestBool(object):
             assert_equal(np.count_nonzero(a), builtins.sum(a.tolist()))
 
 
+class TestZeroSizeFlexible(object):
+    @staticmethod
+    def _zeros(shape, dtype=str):
+        # not constructable directly
+        dtype = np.dtype([('x', dtype, 0)])
+        return np.zeros(shape, dtype=dtype)['x']
+
+    def test_create(self):
+        zs = self._zeros(10, bytes)
+        assert_equal(zs.itemsize, 0)
+        zs = self._zeros(10, np.void)
+        assert_equal(zs.itemsize, 0)
+        zs = self._zeros(10, unicode)
+        assert_equal(zs.itemsize, 0)
+
+    def _test_sort_partition(self, name, kinds, **kwargs):
+        # Previously, these would all hang
+        for dt in [bytes, np.void, unicode]:
+            zs = self._zeros(10, dt)
+            sort_method = getattr(zs, name)
+            sort_func = getattr(np, name)
+            for kind in kinds:
+                sort_method(kind=kind, **kwargs)
+                sort_func(zs, kind=kind, **kwargs)
+
+    def test_sort(self):
+        self._test_sort_partition('sort', kinds='qhm')
+
+    def test_argsort(self):
+        self._test_sort_partition('argsort', kinds='qhm')
+
+    def test_partition(self):
+        self._test_sort_partition('partition', kinds=['introselect'], kth=2)
+
+    def test_argpartition(self):
+        self._test_sort_partition('argpartition', kinds=['introselect'], kth=2)
+
+
 class TestMethods(object):
     def test_compress(self):
         tgt = [[5, 6, 7, 8, 9]]