[Python-checkins] r73531 - in python/branches/py3k: Lib/test/test_memoryview.py Misc/NEWS Objects/memoryobject.c

[raymond.hettinger]

Author: raymond.hettinger
Date: Tue Jun 23 22:38:54 2009
New Revision: 73531

Log:
Issue 6329: Fix iteration for memoryviews.

Modified:
   python/branches/py3k/Lib/test/test_memoryview.py
   python/branches/py3k/Misc/NEWS
   python/branches/py3k/Objects/memoryobject.c

Modified: python/branches/py3k/Lib/test/test_memoryview.py
==============================================================================
--- python/branches/py3k/Lib/test/test_memoryview.py	(original)
+++ python/branches/py3k/Lib/test/test_memoryview.py	Tue Jun 23 22:38:54 2009
@@ -48,6 +48,12 @@
         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

Modified: python/branches/py3k/Misc/NEWS
==============================================================================
--- python/branches/py3k/Misc/NEWS	(original)
+++ python/branches/py3k/Misc/NEWS	Tue Jun 23 22:38:54 2009
@@ -12,6 +12,9 @@
 Core and Builtins
 -----------------
 
+- Issue #6329: Fixed iteration for memoryview objects (it was being blocked
+  because it wasn't recognized as a sequence).
+
 Library
 -------
 

Modified: python/branches/py3k/Objects/memoryobject.c
==============================================================================
--- python/branches/py3k/Objects/memoryobject.c	(original)
+++ python/branches/py3k/Objects/memoryobject.c	Tue Jun 23 22:38:54 2009
@@ -505,6 +505,49 @@
     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
@@ -536,37 +579,7 @@
         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;
@@ -771,6 +784,12 @@
     (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 */
 
@@ -792,7 +811,7 @@
     0,                                        /* tp_reserved */
     (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 */