Merge 023afd7 into b90f413

npctypes/shared.py

@@ -0,0 +1,69 @@
+__author__ = "John Kirkham <kirkhamj@janelia.hhmi.org>"
+__date__ = "$Oct 04, 2016 16:17$"
+
+
+import contextlib
+import multiprocessing
+
+import numpy
+
+from npctypes import types
+
+
+class NDArray(object):
+    def __init__(self, shape, dtype, order=None):
+        if order is None:
+            order = 'C'
+        else:
+            assert order in ['C', 'F'], \
+                "`order` must be `'C'` or `'F'`."
+
+        self._lock = multiprocessing.RLock()
+        self._locked = False
+
+        self._array = multiprocessing.Array(
+            types.ctype(dtype),
+            int(numpy.prod(shape)),
+            lock=False
+        )
+
+        self._type = None
+
+        self._type = numpy.ctypeslib.ndpointer(
+            dtype=dtype,
+            ndim=len(shape),
+            shape=shape,
+            flags=[order]
+        )
+
+    def acquire(self, block=True, timeout=None):
+        return self._lock.acquire(block, timeout)
+
+    def release(self):
+        return self._lock.release()
+
+    @contextlib.contextmanager
+    def as_ndarray(self, lock=True, writeable=True, block=True, timeout=None):
+        if lock:
+            self.acquire(block, timeout)
+
+        # Determine order using bitwise comparison.
+        order = (self._type._flags_ & 1) + (self._type._flags_ & 2)
+        if order == 3:
+            order = 'A'
+        elif order == 2:
+            order = 'F'
+        elif order == 1:
+            order = 'C'
+        else:
+            order = None
+
+        # Construct the NumPy array object.
+        array = numpy.frombuffer(self._array, dtype=self._type._dtype_)
+        array = array.reshape(self._type._shape_, order=order)
+        array.flags["WRITEABLE"] = writeable
+
+        yield array
+
+        if lock:
+            self.release()