From dc5c72d2530e948b9a25138d79d45eb3bd932d76 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Mon, 8 Feb 2021 14:46:29 +0000
Subject: [PATCH] Proof of concept images with units

---
 lib/matplotlib/colorbar.py | 11 +++++++++++
 lib/matplotlib/image.py    | 14 ++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
index 31d9b6176821..fab0a5ec66bd 100644
--- a/lib/matplotlib/colorbar.py
+++ b/lib/matplotlib/colorbar.py
@@ -1174,6 +1174,17 @@ def __init__(self, ax, mappable, **kwargs):
         self.mappable = mappable
         _add_disjoint_kwargs(kwargs, cmap=mappable.cmap, norm=mappable.norm)
 
+        # Set the formatter if
+        # - It isn't explicity specified
+        # - The mappable has a converter and units that provide a formatter
+        if 'format' not in kwargs:
+            if hasattr(mappable, 'converter') and hasattr(mappable, 'units'):
+                info = mappable.converter.axisinfo(mappable.units, self)
+                if info is None:
+                    pass
+                if info.majfmt is not None:
+                    kwargs['format'] = info.majfmt
+
         if isinstance(mappable, contour.ContourSet):
             cs = mappable
             _add_disjoint_kwargs(
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
index 14b3f78f703b..6191e2fb3099 100644
--- a/lib/matplotlib/image.py
+++ b/lib/matplotlib/image.py
@@ -18,6 +18,9 @@
 import matplotlib.colors as mcolors
 import matplotlib.cm as cm
 import matplotlib.cbook as cbook
+# Need to import dates to register unit convertors
+import matplotlib.dates
+import matplotlib.units as munits
 # For clarity, names from _image are given explicitly in this module:
 import matplotlib._image as _image
 # For user convenience, the names from _image are also imported into
@@ -693,6 +696,7 @@ def set_data(self, A):
         """
         if isinstance(A, PIL.Image.Image):
             A = pil_to_array(A)  # Needed e.g. to apply png palette.
+        A = self._convert_units(A)
         self._A = cbook.safe_masked_invalid(A, copy=True)
 
         if (self._A.dtype != np.uint8 and
@@ -730,6 +734,16 @@ def set_data(self, A):
         self._rgbacache = None
         self.stale = True
 
+    def _convert_units(self, A):
+        # Take the first element since units expects a 1D sequence, not 2D
+        converter = munits.registry.get_converter(A[0])
+        if converter is None:
+            return A
+
+        self.converter = converter
+        self.units = converter.default_units(A, self)
+        return converter.convert(A, self.units, self)
+
     def set_array(self, A):
         """
         Retained for backwards compatibility - use set_data instead.