introduction of _ColorbarMappable

trygvrad · trygvrad · commit e418ac7331fa · 2026-01-26T22:09:52.000+01:00
diff --git a/doc/api/colorizer_api.rst b/doc/api/colorizer_api.rst
@@ -6,4 +6,4 @@
    :members:
    :undoc-members:
    :show-inheritance:
-   :private-members: _ColorizerInterface, _ScalarMappable
+   :private-members: _ColorbarMappable, _ScalarMappable
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -1043,6 +1043,8 @@ def remove(self):
 
         try:
             ax = self.mappable.axes
+            if ax is None:
+                return
         except AttributeError:
             return
         try:
diff --git a/lib/matplotlib/colorizer.py b/lib/matplotlib/colorizer.py
@@ -313,7 +313,7 @@ def clip(self, clip):
         self.norm.clip = clip
 
 
-class _ColorizerInterface:
+class _ColorbarMappable:
     """
     Base class that contains the interface to `Colorizer` objects from
     a `ColorizingArtist` or `.cm.ScalarMappable`.
@@ -322,9 +322,94 @@ class _ColorizerInterface:
     attribute. Other functions that as shared between `.ColorizingArtist`
     and `.cm.ScalarMappable` are not included.
     """
+
+    def __init__(self, colorizer, **kwargs):
+        """
+        Base class for objects that can connect to a colorbar.
+
+        All classes that can act as a mappable for `fig.colorbar(mappable)`
+        will subclass this class.
+        """
+        super().__init__(**kwargs)
+        self._colorizer = colorizer
+        self.colorbar = None
+        self._id_colorizer = self._colorizer.callbacks.connect('changed', self.changed)
+        self.callbacks = cbook.CallbackRegistry(signals=["changed"])
+        self._axes = None
+        self._A = None
+
+    @property
+    def axes(self):
+        return self._axes
+
+    @axes.setter
+    def axes(self, axes):
+        self._axes = axes
+
+    @property
+    def colorizer(self):
+        return self._colorizer
+
+    @colorizer.setter
+    def colorizer(self, cl):
+        _api.check_isinstance(Colorizer, colorizer=cl)
+        self._colorizer.callbacks.disconnect(self._id_colorizer)
+        self._colorizer = cl
+        self._id_colorizer = cl.callbacks.connect('changed', self.changed)
+
+    def changed(self):
+        """
+        Call this whenever the mappable is changed to notify all the
+        callbackSM listeners to the 'changed' signal.
+        """
+        self.callbacks.process('changed')
+        self.stale = True
+
     def _scale_norm(self, norm, vmin, vmax):
         self._colorizer._scale_norm(norm, vmin, vmax, self._A)
 
+    def set_array(self, A):
+        """
+        Set the value array from array-like *A*.
+
+        Parameters
+        ----------
+        A : array-like or None
+            The values that are mapped to colors.
+
+            The base class `.ScalarMappable` does not make any assumptions on
+            the dimensionality and shape of the value array *A*.
+        """
+        if A is None:
+            self._A = None
+            return
+
+        A = _ensure_multivariate_data(A, self.norm.n_components)
+
+        A = cbook.safe_masked_invalid(A, copy=True)
+        if not np.can_cast(A.dtype, float, "same_kind"):
+            if A.dtype.fields is None:
+
+                raise TypeError(f"Image data of dtype {A.dtype} cannot be "
+                                f"converted to float")
+            else:
+                for key in A.dtype.fields:
+                    if not np.can_cast(A[key].dtype, float, "same_kind"):
+                        raise TypeError(f"Image data of dtype {A.dtype} cannot be "
+                                        f"converted to a sequence of floats")
+        self._A = A
+        if not self.norm.scaled():
+            self._colorizer.autoscale_None(A)
+
+    def get_array(self):
+        """
+        Return the array of values, that are mapped to colors.
+
+        The base class `.ScalarMappable` does not make any assumptions on
+        the dimensionality and shape of the array.
+        """
+        return self._A
+
     def to_rgba(self, x, alpha=None, bytes=False, norm=True):
         """
         Return a normalized RGBA array corresponding to *x*.
@@ -514,7 +599,7 @@ def _sig_digits_from_norm(norm, data, n):
         return g_sig_digits
 
 
-class _ScalarMappable(_ColorizerInterface):
+class _ScalarMappable(_ColorbarMappable):
     """
     A mixin class to map one or multiple sets of scalar data to RGBA.
 
@@ -527,15 +612,8 @@ class _ScalarMappable(_ColorizerInterface):
     # and ColorizingArtist classes.
 
     # _ScalarMappable can be depreciated so that ColorizingArtist
-    # inherits directly from _ColorizerInterface.
-    # in this case, the following changes should occur:
-    # __init__() has its functionality moved to ColorizingArtist.
-    # set_array(), get_array(), _get_colorizer() and
-    # _check_exclusionary_keywords() are moved to ColorizingArtist.
-    # changed() can be removed so long as colorbar.Colorbar
-    # is changed to connect to the colorizer instead of the
-    # ScalarMappable/ColorizingArtist,
-    # otherwise changed() can be moved to ColorizingArtist.
+    # inherits directly from _ColorbarMappable.
+    # in this case, all functionality should be moved to ColorizingArtist.
     def __init__(self, norm=None, cmap=None, *, colorizer=None, **kwargs):
         """
         Parameters
@@ -550,63 +628,8 @@ def __init__(self, norm=None, cmap=None, *, colorizer=None, **kwargs):
         cmap : str or `~matplotlib.colors.Colormap`
             The colormap used to map normalized data values to RGBA colors.
         """
-        super().__init__(**kwargs)
-        self._A = None
-        self._colorizer = self._get_colorizer(colorizer=colorizer, norm=norm, cmap=cmap)
-
-        self.colorbar = None
-        self._id_colorizer = self._colorizer.callbacks.connect('changed', self.changed)
-        self.callbacks = cbook.CallbackRegistry(signals=["changed"])
-
-    def set_array(self, A):
-        """
-        Set the value array from array-like *A*.
-
-        Parameters
-        ----------
-        A : array-like or None
-            The values that are mapped to colors.
-
-            The base class `.ScalarMappable` does not make any assumptions on
-            the dimensionality and shape of the value array *A*.
-        """
-        if A is None:
-            self._A = None
-            return
-
-        A = _ensure_multivariate_data(A, self.norm.n_components)
-
-        A = cbook.safe_masked_invalid(A, copy=True)
-        if not np.can_cast(A.dtype, float, "same_kind"):
-            if A.dtype.fields is None:
-
-                raise TypeError(f"Image data of dtype {A.dtype} cannot be "
-                                f"converted to float")
-            else:
-                for key in A.dtype.fields:
-                    if not np.can_cast(A[key].dtype, float, "same_kind"):
-                        raise TypeError(f"Image data of dtype {A.dtype} cannot be "
-                                        f"converted to a sequence of floats")
-        self._A = A
-        if not self.norm.scaled():
-            self._colorizer.autoscale_None(A)
-
-    def get_array(self):
-        """
-        Return the array of values, that are mapped to colors.
-
-        The base class `.ScalarMappable` does not make any assumptions on
-        the dimensionality and shape of the array.
-        """
-        return self._A
-
-    def changed(self):
-        """
-        Call this whenever the mappable is changed to notify all the
-        callbackSM listeners to the 'changed' signal.
-        """
-        self.callbacks.process('changed', self)
-        self.stale = True
+        colorizer = self._get_colorizer(colorizer=colorizer, norm=norm, cmap=cmap)
+        super().__init__(colorizer, **kwargs)
 
     @staticmethod
     def _check_exclusionary_keywords(colorizer, **kwargs):
@@ -710,17 +733,6 @@ def __init__(self, colorizer, **kwargs):
         _api.check_isinstance(Colorizer, colorizer=colorizer)
         super().__init__(colorizer=colorizer, **kwargs)
 
-    @property
-    def colorizer(self):
-        return self._colorizer
-
-    @colorizer.setter
-    def colorizer(self, cl):
-        _api.check_isinstance(Colorizer, colorizer=cl)
-        self._colorizer.callbacks.disconnect(self._id_colorizer)
-        self._colorizer = cl
-        self._id_colorizer = cl.callbacks.connect('changed', self.changed)
-
     def _set_colorizer_check_keywords(self, colorizer, **kwargs):
         """
         Raises a ValueError if any kwarg is not None while colorizer is not None.
diff --git a/lib/matplotlib/colorizer.pyi b/lib/matplotlib/colorizer.pyi
@@ -1,4 +1,4 @@
-from matplotlib import cbook, colorbar, colors, artist
+from matplotlib import cbook, colorbar, colors, artist, axes as maxes
 
 import numpy as np
 from numpy.typing import ArrayLike
@@ -46,10 +46,27 @@ class Colorizer:
     def clip(self, value: bool) -> None: ...
 
 
-class _ColorizerInterface:
-    cmap: colors.Colormap
+
+class _ColorbarMappable:
     colorbar: colorbar.Colorbar | None
     callbacks: cbook.CallbackRegistry
+    cmap: colors.Colormap
+    def __init__(
+        self,
+        colorizer: Colorizer | None,
+        **kwargs
+    ) -> None: ...
+    @property
+    def colorizer(self) -> Colorizer: ...
+    @colorizer.setter
+    def colorizer(self, cl: Colorizer) -> None: ...
+    def changed(self) -> None: ...
+    def set_array(self, A: ArrayLike | None) -> None: ...
+    def get_array(self) -> np.ndarray | None: ...
+    @property
+    def axes(self) -> maxes._base._AxesBase | None: ...
+    @axes.setter
+    def axes(self, new_axes: maxes._base._AxesBase | None) -> None: ...
     def to_rgba(
         self,
         x: np.ndarray,
@@ -71,7 +88,7 @@ class _ColorizerInterface:
     def autoscale_None(self) -> None: ...
 
 
-class _ScalarMappable(_ColorizerInterface):
+class _ScalarMappable(_ColorbarMappable):
     def __init__(
         self,
         norm: colors.Norm | None = ...,
@@ -80,9 +97,6 @@ class _ScalarMappable(_ColorizerInterface):
         colorizer: Colorizer | None = ...,
         **kwargs
     ) -> None: ...
-    def set_array(self, A: ArrayLike | None) -> None: ...
-    def get_array(self) -> np.ndarray | None: ...
-    def changed(self) -> None: ...
 
 
 class ColorizingArtist(_ScalarMappable, artist.Artist):
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -78,7 +78,7 @@
 from matplotlib.axes import Subplot  # noqa: F401
 from matplotlib.backends import BackendFilter, backend_registry
 from matplotlib.projections import PolarAxes
-from matplotlib.colorizer import _ColorizerInterface, ColorizingArtist, Colorizer
+from matplotlib.colorizer import _ColorbarMappable, ColorizingArtist, Colorizer
 from matplotlib import mlab  # for detrend_none, window_hanning
 from matplotlib.scale import get_scale_names  # noqa: F401
 
@@ -4201,7 +4201,7 @@ def spy(
         origin=origin,
         **kwargs,
     )
-    if isinstance(__ret, _ColorizerInterface):
+    if isinstance(__ret, _ColorbarMappable):
         sci(__ret)
     return __ret
 
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -15,6 +15,7 @@
     BoundaryNorm, LogNorm, PowerNorm, Normalize, NoNorm
 )
 from matplotlib.colorbar import Colorbar
+from matplotlib.colorizer import Colorizer, _ColorbarMappable
 from matplotlib.ticker import FixedLocator, LogFormatter, StrMethodFormatter
 from matplotlib.testing.decorators import check_figures_equal
 
@@ -1242,3 +1243,24 @@ def test_colorbar_format_string_and_old():
     plt.imshow([[0, 1]])
     cb = plt.colorbar(format="{x}%")
     assert isinstance(cb._formatter, StrMethodFormatter)
+
+
+def test_ColorbarMappable_as_input():
+    # check that _ColorbarMappable can function as a
+    # valid input for colorbar
+    fig, ax = plt.subplots()
+    norm = Normalize(vmin=-5, vmax=10)
+    cl = Colorizer(norm=norm)
+    cbm = _ColorbarMappable(cl)
+    # test without an axes on the mappable, no kwarg
+    with pytest.raises(ValueError, match='Unable to determine Axes to steal'):
+        cb = fig.colorbar(cbm)
+    # test without an axes on the mappable, with kwarg
+    cb = fig.colorbar(cbm, ax=ax)
+    cb = fig.colorbar(cbm, cax=ax)
+    # test without an axes on the mappable
+    cbm.axes = ax
+    cb = fig.colorbar(cbm)
+    assert cb.mappable is cbm
+    assert cb.vmin == -5
+    assert cb.vmax == 10
diff --git a/tools/boilerplate.py b/tools/boilerplate.py
@@ -313,7 +313,7 @@ def boilerplate_gen():
         'hist2d': 'sci(__ret[-1])',
         'imshow': 'sci(__ret)',
         'spy': (
-            'if isinstance(__ret, _ColorizerInterface):\n'
+            'if isinstance(__ret, _ColorbarMappable):\n'
             '        sci(__ret)'
         ),
         'quiver': 'sci(__ret)',
