BUG: Don't segfault on bad __len__ when assigning. (gh-16327)

When __getitem__ fails, assignment falls back on __iter__, which may not
have the same length as __len__, resulting in a segfault.
See gh-7264.

* BUG: Don't rely on __len__ for safe iteration.

Skip __len__ checking entirely, since this has no guarantees of being
correct. Instead, first convert to PySequence_Fast and use that length.
Also fixes a refleak when creating a tmp array fails.
See gh-7264.

* TST: Update test for related edge-case.

The previous fix more gracefully handles this edge case by skipping the
__len__ check. Rather than raising with an unhelpful error message, just
create the array with whatever elements C() actually yields.
See gh-7264.

2 files changed, 25 insertions, 20 deletions

diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index 9283eefce..502ab0ea9 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -483,6 +483,8 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s,
     /* INCREF on entry DECREF on exit */
     Py_INCREF(s);
 
+    PyObject *seq = NULL;
+
     if (PyArray_Check(s)) {
         if (!(PyArray_CheckExact(s))) {
             /*
@@ -529,10 +531,11 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s,
         return 0;
     }
 
-    slen = PySequence_Length(s);
-    if (slen < 0) {
+    seq = PySequence_Fast(s, "Could not convert object to sequence");
+    if (seq == NULL) {
         goto fail;
     }
+    slen = PySequence_Fast_GET_SIZE(seq);
 
     /*
      * Either the dimensions match, or the sequence has length 1 and can
@@ -547,14 +550,9 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s,
 
     /* Broadcast the one element from the sequence to all the outputs */
     if (slen == 1) {
-        PyObject *o;
+        PyObject *o = PySequence_Fast_GET_ITEM(seq, 0);
         npy_intp alen = PyArray_DIM(a, dim);
 
-        o = PySequence_GetItem(s, 0);
-        if (o == NULL) {
-            goto fail;
-        }
-
         for (i = 0; i < alen; i++) {
             if ((PyArray_NDIM(a) - dim) > 1) {
                 PyArrayObject * tmp =
@@ -571,26 +569,18 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s,
                 res = PyArray_SETITEM(dst, b, o);
             }
             if (res < 0) {
-                Py_DECREF(o);
                 goto fail;
             }
         }
-        Py_DECREF(o);
     }
     /* Copy element by element */
     else {
-        PyObject * seq;
-        seq = PySequence_Fast(s, "Could not convert object to sequence");
-        if (seq == NULL) {
-            goto fail;
-        }
         for (i = 0; i < slen; i++) {
             PyObject * o = PySequence_Fast_GET_ITEM(seq, i);
             if ((PyArray_NDIM(a) - dim) > 1) {
                 PyArrayObject * tmp =
                     (PyArrayObject *)array_item_asarray(dst, i);
                 if (tmp == NULL) {
-                    Py_DECREF(seq);
                     goto fail;
                 }
 
@@ -602,17 +592,17 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s,
                 res = PyArray_SETITEM(dst, b, o);
             }
             if (res < 0) {
-                Py_DECREF(seq);
                 goto fail;
             }
         }
-        Py_DECREF(seq);
     }
 
+    Py_DECREF(seq);
     Py_DECREF(s);
     return 0;
 
  fail:
+    Py_XDECREF(seq);
     Py_DECREF(s);
     return res;
 }
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index a698370b6..fc2ab94e6 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -988,7 +988,22 @@ class TestCreation:
             def __len__(self):
                 return 42
 
-        assert_raises(ValueError, np.array, C()) # segfault?
+        a = np.array(C()) # segfault?
+        assert_equal(len(a), 0)
+
+    def test_false_len_iterable(self):
+        # Special case where a bad __getitem__ makes us fall back on __iter__:
+        class C:
+            def __getitem__(self, x):
+                raise Exception
+            def __iter__(self):
+                return iter(())
+            def __len__(self):
+                return 2
+
+        a = np.empty(2)
+        with assert_raises(ValueError):
+            a[:] = C()  # Segfault!
 
     def test_failed_len_sequence(self):
         # gh-7393
@@ -1804,7 +1819,7 @@ class TestMethods:
             c = b.copy()
             c.sort(kind=kind)
             assert_equal(c, a, msg)
-            
+
     def test_sort_structured(self):
         # test record array sorts.
         dt = np.dtype([('f', float), ('i', int)])