Merge pull request #299 from teoliphant/fixup_flat

Fix flat arrays so they cannot be used to write to read-only memory.

3 files changed, 66 insertions, 8 deletions

diff --git a/numpy/core/src/multiarray/getset.c b/numpy/core/src/multiarray/getset.c
index a8470e326..065a207d2 100644
--- a/numpy/core/src/multiarray/getset.c
+++ b/numpy/core/src/multiarray/getset.c
@@ -829,6 +829,7 @@ array_flat_set(PyArrayObject *self, PyObject *val)
                 "Cannot delete array flat iterator");
         return -1;
     }
+    if (PyArray_FailUnlessWriteable(self, "array") < 0) return -1;
     typecode = PyArray_DESCR(self);
     Py_INCREF(typecode);
     arr = (PyArrayObject *)PyArray_FromAny(val, typecode,
diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c
index e8064fefd..0e1a498da 100644
--- a/numpy/core/src/multiarray/iterators.c
+++ b/numpy/core/src/multiarray/iterators.c
@@ -1042,6 +1042,9 @@ iter_ass_subscript(PyArrayIterObject *self, PyObject *ind, PyObject *val)
                 "Cannot delete iterator elements");
         return -1;
     }
+    
+    if (PyArray_FailUnlessWriteable(self->ao, "underlying array") < 0) 
+        return -1;
 
     if (ind == Py_Ellipsis) {
         ind = PySlice_New(NULL, NULL, NULL);
@@ -1224,10 +1227,13 @@ iter_array(PyArrayIterObject *it, PyObject *NPY_UNUSED(op))
 
     /* Two options:
      *  1) underlying array is contiguous
-     *  -- return 1-d wrapper around it
-     * 2) underlying array is not contiguous
-     * -- make new 1-d contiguous array with updateifcopy flag set
-     * to copy back to the old array
+     *     -- return 1-d wrapper around it
+     *  2) underlying array is not contiguous
+     *     -- make new 1-d contiguous array with updateifcopy flag set
+     *        to copy back to the old array
+     * 
+     *  If underlying array is readonly, then we make the output array readonly
+     *     and updateifcopy does not apply. 
      */
     size = PyArray_SIZE(it->ao);
     Py_INCREF(PyArray_DESCR(it->ao));
@@ -1260,10 +1266,15 @@ iter_array(PyArrayIterObject *it, PyObject *NPY_UNUSED(op))
             Py_DECREF(ret);
             return NULL;
         }
-        Py_INCREF(it->ao);
-        if (PyArray_SetUpdateIfCopyBase(ret, it->ao) < 0) {
-            Py_DECREF(ret);
-            return NULL;
+        if (PyArray_ISWRITEABLE(it->ao)) {
+            Py_INCREF(it->ao);
+            if (PyArray_SetUpdateIfCopyBase(ret, it->ao) < 0) {
+                Py_DECREF(ret);
+                return NULL;
+            }
+        }
+        else {
+            PyArray_CLEARFLAGS(ret, NPY_ARRAY_WRITEABLE);
         }
     }
     return ret;
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 056749c19..e9c7a9c4f 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -1673,6 +1673,52 @@ class TestFromBuffer(object):
         yield self.tst_basic, asbytes(''), np.array([]), {}
 
 
+class TestFlat(TestCase):
+    def setUp(self):
+        a0 = arange(20.0)
+        a = a0.reshape(4,5)
+        a0.shape = (4,5)
+        a.flags.writeable = False
+        self.a = a
+        self.b = a[::2,::2]
+        self.a0 = a0
+        self.b0 = a0[::2,::2]
+        
+    def test_contiguous(self):
+        testpassed = False
+        try:
+            self.a.flat[12] = 100.0
+        except ValueError:
+            testpassed = True
+        assert testpassed
+        assert self.a.flat[12] == 12.0
+
+    def test_discontiguous(self):
+        testpassed = False
+        try:
+            self.b.flat[4] = 100.0
+        except ValueError:
+            testpassed = True
+        assert testpassed
+        assert self.b.flat[4] == 12.0
+
+    def test___array__(self):
+        c = self.a.flat.__array__()
+        d = self.b.flat.__array__()
+        e = self.a0.flat.__array__()
+        f = self.b0.flat.__array__()
+
+        assert c.flags.writeable is False
+        assert d.flags.writeable is False
+        assert e.flags.writeable is True
+        assert f.flags.writeable is True
+
+        assert c.flags.updateifcopy is False
+        assert d.flags.updateifcopy is False
+        assert e.flags.updateifcopy is False
+        assert f.flags.updateifcopy is True
+        assert f.base is self.b0    
+
 class TestResize(TestCase):
     def test_basic(self):
         x = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])