Merge pull request #14313 from efiring/masked_to_rgba

dstansby · dstansby · commit 04baaf71b3f1 · 2019-05-25T10:49:27.000+01:00
Support masked array inputs for to_rgba and to_rgba_array.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -152,7 +152,7 @@ def to_rgba(c, alpha=None):
 
     Parameters
     ----------
-    c : Matplotlib color
+    c : Matplotlib color or ``np.ma.masked``
 
     alpha : scalar, optional
         If *alpha* is not ``None``, it forces the alpha value, except if *c* is
@@ -189,6 +189,8 @@ def _to_rgba_no_colorcycle(c, alpha=None):
     ``"none"`` (case-insensitive), which always maps to ``(0, 0, 0, 0)``.
     """
     orig_c = c
+    if c is np.ma.masked:
+        return (0., 0., 0., 0.)
     if isinstance(c, str):
         if c.lower() == "none":
             return (0., 0., 0., 0.)
@@ -254,19 +256,25 @@ def to_rgba_array(c, alpha=None):
 
     If *alpha* is not ``None``, it forces the alpha value.  If *c* is
     ``"none"`` (case-insensitive) or an empty list, an empty array is returned.
+    If *c* is a masked array, an ndarray is returned with a (0, 0, 0, 0)
+    row for each masked value or row in *c*.
     """
     # Special-case inputs that are already arrays, for performance.  (If the
     # array has the wrong kind or shape, raise the error during one-at-a-time
     # conversion.)
     if (isinstance(c, np.ndarray) and c.dtype.kind in "if"
             and c.ndim == 2 and c.shape[1] in [3, 4]):
+        mask = c.mask.any(axis=1) if np.ma.is_masked(c) else None
+        c = np.ma.getdata(c)
         if c.shape[1] == 3:
             result = np.column_stack([c, np.zeros(len(c))])
             result[:, -1] = alpha if alpha is not None else 1.
         elif c.shape[1] == 4:
             result = c.copy()
             if alpha is not None:
                 result[:, -1] = alpha
+        if mask is not None:
+            result[mask] = 0
         if np.any((result < 0) | (result > 1)):
             raise ValueError("RGBA values should be within 0-1 range")
         return result
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -766,6 +766,33 @@ def test_conversions():
         hex_color
 
 
+def test_conversions_masked():
+    x1 = np.ma.array(['k', 'b'], mask=[True, False])
+    x2 = np.ma.array([[0, 0, 0, 1], [0, 0, 1, 1]])
+    x2[0] = np.ma.masked
+    assert mcolors.to_rgba(x1[0]) == (0, 0, 0, 0)
+    assert_array_equal(mcolors.to_rgba_array(x1),
+                       [[0, 0, 0, 0], [0, 0, 1, 1]])
+    assert_array_equal(mcolors.to_rgba_array(x2), mcolors.to_rgba_array(x1))
+
+
+def test_to_rgba_array_single_str():
+    # single color name is valid
+    assert_array_equal(mcolors.to_rgba_array("red"), [(1, 0, 0, 1)])
+
+    # single char color sequence is deprecated
+    with pytest.warns(cbook.MatplotlibDeprecationWarning,
+                      match="Using a string of single character colors as a "
+                            "color sequence is deprecated"):
+        array = mcolors.to_rgba_array("rgb")
+    assert_array_equal(array, [(1, 0, 0, 1), (0, 0.5, 0, 1), (0, 0, 1, 1)])
+
+    with pytest.raises(ValueError,
+                       match="neither a valid single color nor a color "
+                             "sequence"):
+        mcolors.to_rgba_array("rgbx")
+
+
 def test_grey_gray():
     color_mapping = mcolors._colors_full_map
     for k in color_mapping.keys():
