ENH: Make memmaps return arrays in ufuncs and fancy indexing

seberg · seberg · commit d86c8bdd9778 · 2013-01-18T23:27:42.000+01:00
This makes memmeps return an array instead of a memmap instance
for fancy indexing or ufuncs, etc. that use __array_wrap__ and
__array_prepare__. The type is preserved for subclasses however.
Also copy is one example that still returns a memmap even if it
does not point to the same memory.
diff --git a/numpy/core/memmap.py b/numpy/core/memmap.py
@@ -282,6 +282,27 @@ def __array_finalize__(self, obj):
             self.offset = None
             self.mode = None
 
+    def __array_wrap__(self, arr, context=None):
+        # If the type is not memmap, then a user subclass may have some special
+        # functions that should be preserved to avoid a regression
+        if type(self) is not memmap:
+            return arr.view(type=type(self))
+        return arr
+
+    def __array_prepare__(self, arr, context=None):
+        if type(self) is not memmap:
+            return arr.view(type=type(self))
+        # There is no use in making an array a memmap, since it does not own
+        # its memory.
+        return arr
+
+    def __getitem__(self, index):
+        # return an array if the result does not point to self's memory
+        res = super(memmap, self).__getitem__(index)
+        if type(res) is memmap and res._mmap is None:
+            return res.view(ndarray)
+        return res
+
     def flush(self):
         """
         Write any changes in the array to the file on disk.
diff --git a/numpy/core/tests/test_memmap.py b/numpy/core/tests/test_memmap.py
@@ -3,7 +3,7 @@
 import os
 
 from numpy import memmap
-from numpy import arange, allclose, asarray
+from numpy import arange, allclose, asarray, add
 from numpy.testing import *
 
 class TestMemmap(TestCase):
@@ -115,5 +115,20 @@ def test_view(self):
         new_array = asarray(fp)
         assert(new_array.base is fp)
 
+    def test_to_array_conversions(self):
+        class MMSubclass(memmap):
+            pass
+        m = memmap(self.tmpfp, dtype=self.dtype, shape=self.shape)
+        r = m + 1
+        assert_(not isinstance(r, memmap))
+        add(m, 1, out=m)
+        assert_(isinstance(m, memmap))
+        assert_(not isinstance(m[:,[0,1]], memmap))
+        # But for subclasses, preserve in case someone tagged on functions:
+        m = MMSubclass(self.tmpfp, dtype=self.dtype, shape=self.shape)
+        r = m + 1
+        assert_(isinstance(r, MMSubclass))
+        assert_(isinstance(m[:,[0,1]], MMSubclass))
+
 if __name__ == "__main__":
     run_module_suite()
