More work on the memory object and MPI memory allocation

src/MPI/MPI.pyx

@@ -96,13 +96,14 @@ def Alloc_mem(Aint size, Info info=INFO_NULL):
     CHKERR( MPI_Alloc_mem(size, cinfo, &base) )
     return tomemory(base, size)
 
-def Free_mem(memory):
+def Free_mem(mem):
     """
     Free memory allocated with `Alloc_mem()`
     """
     cdef void *base = NULL
-    asmemory(memory, &base, NULL)
+    cdef memory m = asmemory(mem, &base, NULL)
     CHKERR( MPI_Free_mem(base) )
+    m.release()
 
 # Initialization and Exit
 # -----------------------

src/MPI/asbuffer.pxi

@@ -1,5 +1,15 @@
 #------------------------------------------------------------------------------
 
+cdef extern from "Python.h":
+    int PyIndex_Check(object)
+    int PySlice_Check(object)
+    int PySlice_GetIndicesEx(object, Py_ssize_t,
+                             Py_ssize_t *, Py_ssize_t *,
+                             Py_ssize_t *, Py_ssize_t *) except -1
+    Py_ssize_t PyNumber_AsSsize_t(object, object) except -1
+
+#------------------------------------------------------------------------------
+
 # Python 3 buffer interface (PEP 3118)
 cdef extern from "Python.h":
     enum: PY3 "(PY_MAJOR_VERSION>=3)"
@@ -142,6 +152,10 @@ cdef class memory:
 
     cdef Py_buffer view
 
+    def __cinit__(self):
+        PyBuffer_FillInfo(&self.view, <object>NULL,
+                          NULL, 0, 0, PyBUF_SIMPLE)
+
     def __dealloc__(self):
         PyBuffer_Release(&self.view)
 
@@ -162,6 +176,18 @@ cdef class memory:
         def __get__(self):
             return self.view.readonly
 
+    # convenience methods
+
+    def tobytes(self):
+        """Return the data in the buffer as a byte string"""
+        return PyBytes_FromStringAndSize(<char*>self.view.buf, self.view.len)
+
+    def release(self):
+        """Release the underlying buffer exposed by the memory object"""
+        PyBuffer_Release(&self.view)
+        PyBuffer_FillInfo(&self.view, <object>NULL,
+                          NULL, 0, 0, PyBUF_SIMPLE)
+
     # buffer interface (PEP 3118)
 
     def __getbuffer__(self, Py_buffer *view, int flags):
@@ -208,17 +234,35 @@ cdef class memory:
         if i < 0: i += self.view.len
         if i < 0 or i >= self.view.len:
             raise IndexError("index out of range")
-        cdef unsigned char *buf = <unsigned char *>self.view.buf
+        cdef unsigned char *buf = <unsigned char*>self.view.buf
         return <long>buf[i]
 
-    def __setitem__(self, Py_ssize_t i, unsigned char v):
+    def __setitem__(self, object item, object value):
         if self.view.readonly:
             raise TypeError("memory buffer is read-only")
-        if i < 0: i += self.view.len
-        if i < 0 or i >= self.view.len:
-            raise IndexError("index out of range")
         cdef unsigned char *buf = <unsigned char*>self.view.buf
-        buf[i] = v
+        cdef Py_ssize_t start, stop, step, length
+        cdef memory inmem
+        if PyIndex_Check(item):
+            start = PyNumber_AsSsize_t(item, IndexError)
+            if start < 0: start += self.view.len
+            if start < 0 or start >= self.view.len:
+                raise IndexError("index out of range")
+            buf[start] = <unsigned char>value
+        elif PySlice_Check(item):
+            PySlice_GetIndicesEx(item, self.view.len,
+                                 &start, &stop, &step, &length)
+            if step != 1:
+                raise IndexError("slice with step not supported")
+            if PyIndex_Check(value):
+                <void>memset(buf+start, <unsigned char>value, <size_t>length)
+            else:
+                inmem = getbuffer(value, 1, 0)
+                if inmem.view.len != length:
+                    raise ValueError("slice length does not match buffer")
+                <void>memcpy(buf+start, inmem.view.buf, <size_t>length)
+        else:
+            raise TypeError("indices must be integers or slices")
 
 #------------------------------------------------------------------------------
 
@@ -270,23 +314,29 @@ cdef inline object getformat(memory buf):
 
 cdef inline memory getbuffer_r(object ob, void **base, MPI_Aint *size):
     cdef memory buf = getbuffer(ob, 1, 0)
-    if base != NULL: base[0] = <void*>    buf.view.buf
-    if size != NULL: size[0] = <MPI_Aint> buf.view.len
+    if base != NULL: base[0] = buf.view.buf
+    if size != NULL: size[0] = <MPI_Aint>buf.view.len
     return buf
 
 cdef inline memory getbuffer_w(object ob, void **base, MPI_Aint *size):
     cdef memory buf = getbuffer(ob, 0, 0)
-    if base != NULL: base[0] = <void*>    buf.view.buf
-    if size != NULL: size[0] = <MPI_Aint> buf.view.len
+    if base != NULL: base[0] = buf.view.buf
+    if size != NULL: size[0] = <MPI_Aint>buf.view.len
     return buf
 
 #------------------------------------------------------------------------------
 
-cdef inline object asmemory(object ob, void **base, MPI_Aint *size):
-    cdef memory buf = getbuffer_w(ob, base, size)
-    return buf
+cdef inline memory asmemory(object ob, void **base, MPI_Aint *size):
+    cdef memory mem
+    if isinstance(ob, memory):
+        mem = <memory> ob
+    else:
+        mem = getbuffer(ob, 1, 0)
+    if base != NULL: base[0] = mem.view.buf
+    if size != NULL: size[0] = <MPI_Aint> mem.view.len
+    return mem
 
-cdef inline object tomemory(void *base, MPI_Aint size):
+cdef inline memory tomemory(void *base, MPI_Aint size):
     return tobuffer(base, size)
 
 #------------------------------------------------------------------------------

src/MPI/stdlib.pxi

@@ -1,5 +1,10 @@
 #------------------------------------------------------------------------------
 
+cdef extern from * nogil: # "string.h"
+    void *memset(void *, int, size_t)
+    void *memcpy(void *, void *, size_t)
+    void *memmove(void *, void *, size_t)
+
 cdef extern from * nogil: # "stdio.h"
     ctypedef struct FILE
     FILE *stdin, *stdout, *stderr

src/pycompat.h

@@ -30,6 +30,24 @@ _PyLong_AsByteArray(PyLongObject* v,
 
 /* ------------------------------------------------------------------------- */
 
+#ifdef PYPY_VERSION
+#ifdef PySlice_GetIndicesEx
+#undef PySlice_GetIndicesEx
+#define PySlice_GetIndicesEx(s, n, start, stop, step, length) \
+PyPySlice_GetIndicesEx((PySliceObject *)(s), n, start, stop, step, length)
+#else
+#define PySlice_GetIndicesEx(s, n, start, stop, step, length) \
+PySlice_GetIndicesEx((PySliceObject *)(s), n, start, stop, step, length)
+#endif
+#else
+#if PY_VERSION_HEX < 0x03020000
+#define PySlice_GetIndicesEx(s, n, start, stop, step, length) \
+PySlice_GetIndicesEx((PySliceObject *)(s), n, start, stop, step, length)
+#endif
+#endif
+
+/* ------------------------------------------------------------------------- */
+
 #if !defined(WITH_THREAD)
 #undef  PyGILState_Ensure
 #define PyGILState_Ensure() ((PyGILState_STATE)0)

test/test_mem.py

@@ -23,5 +23,65 @@ def testMemory2(self):
             except NotImplementedError:
                 return
 
+    def testMemorySequence(self):
+        n = 16
+        try:
+            mem = MPI.Alloc_mem(n, MPI.INFO_NULL)
+        except NotImplementedError:
+            return
+        try:
+            mem.address
+        except AttributeError:
+            MPI.Free_mem(mem)
+            return
+        try:
+            self.assertEqual(len(mem), n)
+            self.assertTrue(mem.address != 0)
+            self.assertEqual(mem.nbytes, n)
+            self.assertFalse(mem.readonly)
+            def getitem(): return mem[n]
+            def delitem(): del mem[n]
+            def setitem1(): mem[n] = 0
+            def setitem2(): mem[::2] = 0
+            def setitem3(): mem[None] = 0
+            self.assertRaises(IndexError, getitem)
+            self.assertRaises(Exception,  delitem)
+            self.assertRaises(IndexError, setitem1)
+            self.assertRaises(IndexError, setitem2)
+            self.assertRaises(TypeError,  setitem3)
+            for i in range(n):
+                mem[i] = i
+            for i in range(n):
+                self.assertEqual(mem[i], i)
+            mem[:] = 0
+            for i in range(n):
+                self.assertEqual(mem[i], 0)
+            mem[:] = 255
+            for i in range(n):
+                self.assertEqual(mem[i], 255)
+            mem[:n//2] = 1
+            mem[n//2:] = 0
+            for i in range(n//2):
+                self.assertEqual(mem[i], 1)
+            for i in range(n//2, n):
+                self.assertEqual(mem[i], 0)
+            mem[:] = 0
+            mem[1:5] = b"abcd"
+            mem[10:13] = b"xyz"
+            self.assertEqual(mem[0], 0)
+            for i, c in enumerate("abcd"):
+                self.assertEqual(mem[1+i], ord(c))
+            for i in range(5, 10):
+                self.assertEqual(mem[i], 0)
+            for i, c in enumerate("xyz"):
+                self.assertEqual(mem[10+i], ord(c))
+            for i in range(13, n):
+                self.assertEqual(mem[i], 0)
+        finally:
+            MPI.Free_mem(mem)
+            self.assertEqual(mem.address, 0)
+            self.assertEqual(mem.nbytes, 0)
+            self.assertFalse(mem.readonly)
+
 if __name__ == '__main__':
     unittest.main()