fix for 2D numpy arrays

mplotutils/_colorbar.py

@@ -193,8 +193,7 @@ def colorbar(
     """
 
     ax = _deprecate_ax1_ax2(ax, ax2, kwargs.pop("ax1", None))
-
-    axs = [ax] if np.ndim(ax) == 0 else ax
+    axs = np.asarray(ax).flatten()
 
     if orientation not in ("vertical", "horizontal"):
         raise ValueError("orientation must be 'vertical' or 'horizontal'")

mplotutils/tests/test_colorbar.py

@@ -155,15 +155,6 @@ def test_colorbar_deprecate_ax1():
         with pytest.warns(match="`ax1` has been deprecated in favor of `ax`"):
             mpu.colorbar(h, ax1=ax)
 
-    # create_figure_subplots
-
-    # # with pytest.warns(match=""):
-    # mpu.colorbar(h, ax, ax)
-
-    # mpu.colorbar(h, ax, ax2=ax, orientation=orientation)
-
-    # mpu.colorbar(h, ax1=ax, ax2=ax, orientation=orientation)
-
 
 def test_colorbar_different_figures():
     with figure_context() as f1, figure_context() as f2:
@@ -221,6 +212,29 @@ def colorbar_one_ax_horizontal(**kwargs):
     return cbar
 
 
+def test_colorbar_2d_array():
+    # ensure passing a 2x2 numpy array of axes works
+
+    with figure_context():
+        h, axs = create_figure_subplots(2, 2)
+
+        cbar = mpu.colorbar(h, axs, size=0.2, pad=0)
+
+        expected = [0.8, 0, 0.2 * 0.8, 1.0]
+        assert_position(cbar, expected)
+
+    with figure_context():
+        h, axs = create_figure_subplots(2, 2, orientation="horizontal")
+
+        cbar = mpu.colorbar(
+            h, [axs[0], axs[1]], size=0.2, pad=0, orientation="horizontal"
+        )
+
+        height = 0.2 * 0.8
+        expected = [0.0, 0.2 - height, 1.0, height]
+        assert_position(cbar, expected)
+
+
 def test_colorbar_vertical_aspect():
     with figure_context(figsize=(5, 5)):
         # test pad=0, aspect=5