fix: treat unmasked nans as bad when the input is a masked array

src/cmap/_colormap.py

@@ -405,8 +405,14 @@ def __call__(
 
         mask_under = xa < 0
         mask_over = xa >= N
-        # If input was masked, get the bad mask from it; else mask out nans.
-        mask_bad = x.mask if np.ma.is_masked(x) else np.isnan(xa)  # type: ignore
+        # If input was masked, start from its mask: a masked array can still carry
+        # unmasked nans.  `|` rather than `|=`, so x's own mask isn't written to.
+        if np.ma.is_masked(x):
+            mask_bad = x.mask  # type: ignore
+            if xa.dtype.kind == "f":
+                mask_bad = mask_bad | np.isnan(xa)
+        else:
+            mask_bad = np.isnan(xa)
 
         with np.errstate(invalid="ignore"):
             # We need this cast for unsigned ints as well as floats

tests/test_colormap.py

@@ -144,6 +144,23 @@ def test_colormap_apply() -> None:
     assert cmap1(swapped).shape == (10, 10, 4)
 
 
+def test_colormap_masked_array_with_unmasked_nan() -> None:
+    cmap = Colormap("viridis", bad="red")
+    mask = [True, False, False, False]
+    data = np.ma.masked_array([0.0, 0.25, np.nan, 1.0], mask=mask)
+
+    rgba = cmap(data)
+
+    npt.assert_array_equal(rgba[0], Color("red").rgba)  # masked
+    npt.assert_array_equal(rgba[2], Color("red").rgba)  # nan, not masked
+    npt.assert_array_equal(rgba[[1, 3]], cmap(np.array([0.25, 1.0])))
+    npt.assert_array_equal(data.mask, mask)
+
+    # an all-false mask takes the same path as a plain array
+    all_false = np.ma.masked_array([0.25, np.nan], mask=[False, False])
+    npt.assert_array_equal(cmap(all_false), cmap(np.array([0.25, np.nan])))
+
+
 def test_fill_stops() -> None:
     assert _fill_stops([None, None, None]) == [0, 0.5, 1.0]
     assert _fill_stops([None, 0.8, None]) == [0, 0.8, 1.0]