Merge pull request #20788 from greglucas/cbar-2d-alpha

tacaswell · web-flow · commit a0bcd8c93eb0 · 2021-08-06T12:33:01.000-04:00
FIX: Check for colorbar creation with multi-dimensional alpha
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -427,7 +427,9 @@ def __init__(self, ax, mappable=None, *, cmap=None,
                 extend = norm.extend
             else:
                 extend = 'neither'
-        self.alpha = alpha
+        self.alpha = None
+        # Call set_alpha to handle array-like alphas properly
+        self.set_alpha(alpha)
         self.cmap = cmap
         self.norm = norm
         self.values = values
@@ -934,8 +936,13 @@ def set_label(self, label, *, loc=None, **kwargs):
         self.stale = True
 
     def set_alpha(self, alpha):
-        """Set the transparency between 0 (transparent) and 1 (opaque)."""
-        self.alpha = alpha
+        """
+        Set the transparency between 0 (transparent) and 1 (opaque).
+
+        If an array is provided, *alpha* will be set to None to use the
+        transparency values associated with the colormap.
+        """
+        self.alpha = None if isinstance(alpha, np.ndarray) else alpha
 
     def _set_scale(self, scale, **kwargs):
         """
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -604,6 +604,18 @@ def test_mappable_no_alpha():
     plt.draw()
 
 
+def test_mappable_2d_alpha():
+    fig, ax = plt.subplots()
+    x = np.arange(1, 5).reshape(2, 2)/4
+    pc = ax.pcolormesh(x, alpha=x)
+    cb = fig.colorbar(pc, ax=ax)
+    # The colorbar's alpha should be None and the mappable should still have
+    # the original alpha array
+    assert cb.alpha is None
+    assert pc.get_alpha() is x
+    fig.draw_no_output()
+
+
 def test_colorbar_label():
     """
     Test the label parameter. It should just be mapped to the xlabel/ylabel of
