From e6167803b9c67902997faaec5906db7d1e500f5f Mon Sep 17 00:00:00 2001
From: Greg Lucas <greg.lucas@lasp.colorado.edu>
Date: Tue, 3 Aug 2021 19:38:59 -0600
Subject: [PATCH] FIX: Check for colorbar creation with multi-dimensional alpha

If an array-like alpha is provided, we set the colorbar alpha to None
to use the colormap's transparency values. This would previously error.
---
 lib/matplotlib/colorbar.py            | 13 ++++++++++---
 lib/matplotlib/tests/test_colorbar.py | 12 ++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
index 0b8953d5241f..36dfcbdecf72 100644
--- 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
index b0641f1635b8..055c4acb7642 100644
--- 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