diff options
| -rw-r--r-- | numpy/core/src/multiarray/arrayobject.c | 9 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/ctors.c | 54 | ||||
| -rw-r--r-- | numpy/core/tests/test_regression.py | 16 |
3 files changed, 65 insertions, 14 deletions
diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c index ba4b2fc59..4557bd94d 100644 --- a/numpy/core/src/multiarray/arrayobject.c +++ b/numpy/core/src/multiarray/arrayobject.c @@ -144,10 +144,11 @@ PyArray_CopyObject(PyArrayObject *dest, PyObject *src_object) } } else { - /* If the dims match exactly, can assign directly */ - if (ndim == PyArray_NDIM(dest) && - PyArray_CompareLists(dims, PyArray_DIMS(dest), - ndim)) { + /* + * If there are more than enough dims, use AssignFromSequence + * because it can handle this style of broadcasting. + */ + if (ndim >= PyArray_NDIM(dest)) { int res; Py_DECREF(dtype); res = PyArray_AssignFromSequence(dest, src_object); diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c index 8d4a7430b..275177559 100644 --- a/numpy/core/src/multiarray/ctors.c +++ b/numpy/core/src/multiarray/ctors.c @@ -542,26 +542,60 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s, int dim, npy_intp offset) if (slen < 0) { goto fail; } - if (slen != a->dimensions[dim]) { + /* + * Either the dimensions match, or the sequence has length 1 and can + * be broadcast to the destination. + */ + if (slen != a->dimensions[dim] && slen != 1) { PyErr_Format(PyExc_ValueError, "cannot copy sequence with size %d to array axis " "with dimension %d", (int)slen, (int)a->dimensions[dim]); goto fail; } - for (i = 0; i < slen; i++) { - PyObject *o = PySequence_GetItem(s, i); - if ((a->nd - dim) > 1) { - res = setArrayFromSequence(a, o, dim+1, offset); + /* Broadcast the one element from the sequence to all the outputs */ + if (slen == 1) { + PyObject *o; + npy_intp alen = a->dimensions[dim]; + + o = PySequence_GetItem(s, 0); + if (o == NULL) { + goto fail; } - else { - res = a->descr->f->setitem(o, (a->data + offset), a); + for (i = 0; i < alen; i++) { + if ((a->nd - dim) > 1) { + res = setArrayFromSequence(a, o, dim+1, offset); + } + else { + res = a->descr->f->setitem(o, (a->data + offset), a); + } + if (res < 0) { + Py_DECREF(o); + goto fail; + } + offset += a->strides[dim]; } Py_DECREF(o); - if (res < 0) { - goto fail; + } + /* Copy element by element */ + else { + for (i = 0; i < slen; i++) { + PyObject *o = PySequence_GetItem(s, i); + if (o == NULL) { + goto fail; + } + if ((a->nd - dim) > 1) { + res = setArrayFromSequence(a, o, dim+1, offset); + } + else { + res = a->descr->f->setitem(o, (a->data + offset), a); + } + Py_DECREF(o); + if (res < 0) { + goto fail; + } + offset += a->strides[dim]; } - offset += a->strides[dim]; } Py_DECREF(s); diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py index 5db874ecd..a4cc2d725 100644 --- a/numpy/core/tests/test_regression.py +++ b/numpy/core/tests/test_regression.py @@ -1613,5 +1613,21 @@ class TestRegression(TestCase): b[...] = a_obj assert_equal(b, a_rec) + def test_assign_obj_listoflists(self): + # Ticket # 1870 + # The inner list should get assigned to the object elements + a = np.zeros(4, dtype=object) + b = a.copy() + a[0] = [1] + a[1] = [2] + a[2] = [3] + a[3] = [4] + b[...] = [[1], [2], [3], [4]] + assert_equal(a, b) + # The first dimension should get broadcast + a = np.zeros((2,2), dtype=object) + a[...] = [[1,2]] + assert_equal(a, [[1,2], [1,2]]) + if __name__ == "__main__": run_module_suite() |