mrecords : Make sure a field shares its mask with the whole array

mrecords : IMPORTANT : the mask of a field is no longer set to nomask when it's full of False, which simplifies masking specific fields.
extras   : Reorganized personal comments

4 files changed, 34 insertions, 22 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 5c38ff92b..747fb032f 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -1633,7 +1633,7 @@ class MaskedArray(numeric.ndarray):
                     else:
                         return str(self._data)
                 # convert to object array to make filled work
-#CHECK: the two lines below seem more robust than the self._data.astype
+#!!!: the two lines below seem more robust than the self._data.astype
 #                res = numeric.empty(self._data.shape, object_)
 #                numeric.putmask(res,~m,self._data)
                 res = self._data.astype("|O8")
@@ -2032,7 +2032,7 @@ masked_%(name)s(data = %(data)s,
         indicated `axis1` and `axis2`.
 
         """
-        # TODO: What are we doing with `out`?
+        #!!!: implement out + test!
         m = self._mask
         if m is nomask:
             result = super(MaskedArray, self).trace(offset=offset, axis1=axis1,
@@ -2207,7 +2207,7 @@ masked_%(name)s(data = %(data)s,
 
         """
         if self._mask is nomask:
-            # TODO: Do we keep super, or var _data and take a view ?
+            #???: Do we keep super, or var _data and take a view ?
             return super(MaskedArray, self).var(axis=axis, dtype=dtype,
                                                 ddof=ddof)
         else:
@@ -2676,7 +2676,7 @@ def array(data, dtype=None, copy=False, order=False,
     for convenience.  And backwards compatibility...
 
     """
-    #TODO: we should try to put 'order' somwehere
+    #!!!: we should try to put 'order' somwehere
     return MaskedArray(data, mask=mask, dtype=dtype, copy=copy, subok=subok,
                        keep_mask=keep_mask, hard_mask=hard_mask,
                        fill_value=fill_value, ndmin=ndmin, shrink=shrink)
@@ -3186,7 +3186,7 @@ def where (condition, x=None, y=None):
 
 def choose (indices, t, out=None, mode='raise'):
     "Return array shaped like indices with elements chosen from t"
-    #TODO: implement options `out` and `mode`, if possible.
+    #!!!: implement options `out` and `mode`, if possible + test.
     def fmask (x):
         "Returns the filled array, or True if masked."
         if x is masked:
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index 3032a462b..2ab2eae93 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -567,7 +567,7 @@ def dot(a,b, strict=False):
         The first argument is not conjugated.
 
     """
-    #TODO: Works only with 2D arrays. There should be a way to get it to run with higher dimension
+    #!!!: Works only with 2D arrays. There should be a way to get it to run with higher dimension
     if strict and (a.ndim == 2) and (b.ndim == 2):
         a = mask_rows(a)
         b = mask_cols(b)
diff --git a/numpy/ma/mrecords.py b/numpy/ma/mrecords.py
index 55551fbd0..8ce36c9ad 100644
--- a/numpy/ma/mrecords.py
+++ b/numpy/ma/mrecords.py
@@ -6,11 +6,11 @@ By comparison, mrecarrays support masking individual fields.
 
 :author: Pierre Gerard-Marchant
 """
-#TODO: We should make sure that no field is called '_mask','mask','_fieldmask',
-#TODO: ...or whatever restricted keywords.
-#TODO: An idea would be to no bother in the first place, and then rename the
-#TODO: invalid fields with a trailing underscore...
-#TODO: Maybe we could just overload the parser function ?
+#!!!: * We should make sure that no field is called '_mask','mask','_fieldmask',
+#!!!:   or whatever restricted keywords.
+#!!!:   An idea would be to no bother in the first place, and then rename the
+#!!!:   invalid fields with a trailing underscore...
+#!!!:   Maybe we could just overload the parser function ?
 
 
 __author__ = "Pierre GF Gerard-Marchant"
@@ -51,9 +51,6 @@ def _getformats(data):
     formats = ''
     for obj in data:
         obj = np.asarray(obj)
-#        if not isinstance(obj, ndarray):
-##        if not isinstance(obj, ndarray):
-#            raise ValueError, "item in the array list must be an ndarray."
         formats += _typestr[obj.dtype.type]
         if issubclass(obj.dtype.type, ntypes.flexible):
             formats += `obj.itemsize`
@@ -124,7 +121,6 @@ class MaskedRecords(MaskedArray, object):
         self = recarray.__new__(cls, shape, dtype=dtype, buf=buf, offset=offset,
                                 strides=strides, formats=formats,
                                 byteorder=byteorder, aligned=aligned,)
-#        self = self.view(cls)
         #
         mdtype = [(k,'|b1') for (k,_) in self.dtype.descr]
         if mask is nomask or not np.size(mask):
@@ -331,11 +327,13 @@ The fieldname base is either `_data` or `_mask`."""
         _data = self._data
         # We want a field ........
         if isinstance(indx, basestring):
+            #NB: Make sure _sharedmask is True to propagate back to _fieldmask
+            #NB: Don't use _set_mask, there are some copies being made that break propagation
+            #NB: Don't force the mask to nomask, that wrecks easy masking
             obj = _data[indx].view(MaskedArray)
-            obj._set_mask(_fieldmask[indx])
-            # Force to nomask if the mask is empty
-            if not obj._mask.any():
-                obj._mask = nomask
+#            obj._set_mask(_fieldmask[indx])
+            obj._mask = _fieldmask[indx]
+            obj._sharedmask = True
             # Force to masked if the mask is True
             if not obj.ndim and obj._mask:
                 return masked
@@ -780,4 +778,3 @@ set to 'fi', where `i` is the number of existing fields.
     newdata._fieldmask = newmask
     return newdata
 
-###############################################################################
diff --git a/numpy/ma/tests/test_mrecords.py b/numpy/ma/tests/test_mrecords.py
index b58d1af8d..0d362d3bf 100644
--- a/numpy/ma/tests/test_mrecords.py
+++ b/numpy/ma/tests/test_mrecords.py
@@ -132,6 +132,20 @@ class TestMRecords(NumpyTestCase):
         assert_equal(rdata.num, val)
         assert_equal(rdata.num.mask, [1,0,0])
 
+    def test_set_fields_mask(self):
+        "Tests setting the mask of a field."
+        base = self.base.copy()
+        # This one has already a mask....
+        mbase = base.view(mrecarray)
+        mbase['a'][-2] = masked
+        assert_equal(mbase.a, [1,2,3,4,5])
+        assert_equal(mbase.a._mask, [0,1,0,1,1])
+        # This one has not yet
+        mbase = fromarrays([np.arange(5), np.random.rand(5)],
+                           dtype=[('a',int),('b',float)])
+        mbase['a'][-2] = masked
+        assert_equal(mbase.a, [0,1,2,3,4])
+        assert_equal(mbase.a._mask, [0,0,0,1,0])
     #
     def test_set_mask(self):
         base = self.base.copy()
@@ -231,7 +245,8 @@ class TestMRecords(NumpyTestCase):
         mbase.soften_mask()
         assert(not mbase._hardmask)
         mbase._mask = nomask
-        assert(mbase['b']._mask is nomask)
+        # So, the mask of a field is no longer set to nomask...
+        assert(ma.make_mask(mbase['b']._mask) is nomask)
         assert_equal(mbase['a']._mask,mbase['b']._mask)
     #
     def test_pickling(self):
@@ -365,7 +380,7 @@ class TestMRecordsImport(NumpyTestCase):
         f.write(fcontent)
         f.close()
         mrectxt = fromtextfile(fname,delimitor=',',varnames='ABCDEFG')
-        os.unlink(fname)
+        os.remove(fname)
         #
         assert(isinstance(mrectxt, MaskedRecords))
         assert_equal(mrectxt.F, [1,1,1,1])