Issue 6329: Fix iteration for memoryviews.

rhettinger · rhettinger · commit 62641e9534d7 · 2009-06-23T20:59:43.000Z
diff --git a/Lib/test/test_memoryview.py b/Lib/test/test_memoryview.py
@@ -48,6 +48,12 @@ def test_getitem(self):
         for tp in self._types:
             self.check_getitem_with_type(tp)
 
+    def test_iter(self):
+        for tp in self._types:
+            b = tp(self._source)
+            m = self._view(b)
+            self.assertEqual(list(m), [m[i] for i in range(len(m))])
+
     def test_setitem_readonly(self):
         if not self.ro_type:
             return
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -12,6 +12,9 @@ What's New in Python 2.7 alpha 1
 Core and Builtins
 -----------------
 
+- Issue #6329: Fixed iteration for memoryview objects (it was being blocked
+  because it wasn't recognized as a sequence).
+
 - Issue #6289: Encoding errors from compile() were being masked.
 
 - When no module is given in a relative import, the module field of the
diff --git a/Objects/memoryobject.c b/Objects/memoryobject.c
@@ -513,6 +513,49 @@ memory_length(PyMemoryViewObject *self)
     return get_shape0(&self->view);
 }
 
+/* Alternate version of memory_subcript that only accepts indices.
+   Used by PySeqIter_New().
+*/
+static PyObject *
+memory_item(PyMemoryViewObject *self, Py_ssize_t result)
+{
+    Py_buffer *view = &(self->view);
+
+    if (view->ndim == 0) {
+        PyErr_SetString(PyExc_IndexError,
+                        "invalid indexing of 0-dim memory");
+        return NULL;
+    }
+    if (view->ndim == 1) {
+        /* Return a bytes object */
+        char *ptr;
+        ptr = (char *)view->buf;
+        if (result < 0) {
+            result += get_shape0(view);
+        }
+        if ((result < 0) || (result >= get_shape0(view))) {
+            PyErr_SetString(PyExc_IndexError,
+                                "index out of bounds");
+            return NULL;
+        }
+        if (view->strides == NULL)
+            ptr += view->itemsize * result;
+        else
+            ptr += view->strides[0] * result;
+        if (view->suboffsets != NULL &&
+            view->suboffsets[0] >= 0) {
+            ptr = *((char **)ptr) + view->suboffsets[0];
+        }
+        return PyBytes_FromStringAndSize(ptr, view->itemsize);
+    } else {
+        /* Return a new memory-view object */
+        Py_buffer newview;
+        memset(&newview, 0, sizeof(newview));
+        /* XXX:  This needs to be fixed so it actually returns a sub-view */
+        return PyMemoryView_FromBuffer(&newview);
+    }
+}
+
 /*
   mem[obj] returns a bytes object holding the data for one element if
            obj fully indexes the memory view or another memory-view object
@@ -544,37 +587,7 @@ memory_subscript(PyMemoryViewObject *self, PyObject *key)
         result = PyNumber_AsSsize_t(key, NULL);
         if (result == -1 && PyErr_Occurred())
                 return NULL;
-        if (view->ndim == 1) {
-            /* Return a bytes object */
-            char *ptr;
-            ptr = (char *)view->buf;
-            if (result < 0) {
-                result += get_shape0(view);
-            }
-            if ((result < 0) || (result >= get_shape0(view))) {
-                PyErr_SetString(PyExc_IndexError,
-                                "index out of bounds");
-                return NULL;
-            }
-            if (view->strides == NULL)
-                ptr += view->itemsize * result;
-            else
-                ptr += view->strides[0] * result;
-            if (view->suboffsets != NULL &&
-                view->suboffsets[0] >= 0) {
-                ptr = *((char **)ptr) + view->suboffsets[0];
-            }
-            return PyBytes_FromStringAndSize(ptr, view->itemsize);
-        }
-        else {
-            /* Return a new memory-view object */
-            Py_buffer newview;
-            memset(&newview, 0, sizeof(newview));
-            /* XXX:  This needs to be fixed so it
-                         actually returns a sub-view
-            */
-            return PyMemoryView_FromBuffer(&newview);
-        }
+        return memory_item(self, result);
     }
     else if (PySlice_Check(key)) {
         Py_ssize_t start, stop, step, slicelength;
@@ -779,6 +792,12 @@ static PyMappingMethods memory_as_mapping = {
     (objobjargproc)memory_ass_sub,        /* mp_ass_subscript */
 };
 
+static PySequenceMethods memory_as_sequence = {
+	0,                                  /* sq_length */
+	0,                                  /* sq_concat */
+	0,                                  /* sq_repeat */
+	(ssizeargfunc)memory_item,          /* sq_item */
+};
 
 /* Buffer methods */
 static PyBufferProcs memory_as_buffer = {
@@ -803,7 +822,7 @@ PyTypeObject PyMemoryView_Type = {
     0,                                        /* tp_compare */
     (reprfunc)memory_repr,                    /* tp_repr */
     0,                                        /* tp_as_number */
-    0,                                        /* tp_as_sequence */
+    &memory_as_sequence,                      /* tp_as_sequence */
     &memory_as_mapping,                       /* tp_as_mapping */
     0,                                        /* tp_hash */
     0,                                        /* tp_call */
