MNT: make fig.colorbar(..., ax=INPUT) even more forgiving

tacaswell · tacaswell · commit 9c36d85feaa8 · 2022-12-15T12:21:08.000-05:00
Accept anything that passes `np.iterable`
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -11,7 +11,6 @@
    End-users most likely won't need to directly use this module's API.
 """
 
-import collections.abc as collections_abc
 import logging
 
 import numpy as np
@@ -1395,7 +1394,7 @@ def make_axes(parents, location=None, orientation=None, fraction=0.15,
 
     Parameters
     ----------
-    parents : `~.axes.Axes` or sequence or `numpy.ndarray` of `~.axes.Axes`
+    parents : `~.axes.Axes` or iterable of `~.axes.Axes`
         The Axes to use as parents for placing the colorbar.
     %(_make_axes_kw_doc)s
 
@@ -1419,10 +1418,14 @@ def make_axes(parents, location=None, orientation=None, fraction=0.15,
     # turn parents into a list if it is not already.  Note we cannot
     # use .flatten or .ravel as these copy the references rather than
     # reuse them, leading to a memory leak
-    if isinstance(parents, np.ndarray):
-        parents = list(parents.flat)
-    elif not isinstance(parents, collections_abc.Sequence):
+    if np.iterable(parents):
+        if isinstance(parents, np.ndarray):
+            parents = list(parents.flat)
+        else:
+            parents = list(parents)
+    else:
         parents = [parents]
+
     fig = parents[0].get_figure()
 
     pad0 = 0.05 if fig.get_constrained_layout() else loc_settings['pad']
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -1195,7 +1195,7 @@ def colorbar(
         cax : `~matplotlib.axes.Axes`, optional
             Axes into which the colorbar will be drawn.
 
-        ax : `~.axes.Axes` or sequence or `numpy.ndarray` of Axes, optional
+        ax : `~.axes.Axes` or iterable of Axes, optional
             One or more parent axes from which space for a new colorbar axes
             will be stolen, if *cax* is None.  This has no effect if *cax* is
             set.
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -1209,4 +1209,5 @@ def test_colorbar_axes_parmeters():
     fig.colorbar(im, ax=ax[0])
     fig.colorbar(im, ax=[_ax for _ax in ax])
     fig.colorbar(im, ax=(ax[0], ax[1]))
+    fig.colorbar(im, ax={i: _ax for i, _ax in enumerate(ax)}.values())
     fig.draw_without_rendering()
