Also convert datasets (optionally) in _process_unit_info.

anntzer · anntzer · commit b543bc0e58fb · 2020-09-17T20:13:07.000+02:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -839,8 +839,7 @@ def axhline(self, y=0, xmin=0, xmax=1, **kwargs):
         ymin, ymax = self.get_ybound()
 
         # Strip away the units for comparison with non-unitized bounds.
-        self._process_unit_info({"y": y}, kwargs)
-        yy = self.convert_yunits(y)
+        yy, = self._process_unit_info({"y": y}, kwargs)
         scaley = (yy < ymin) or (yy > ymax)
 
         trans = self.get_yaxis_transform(which='grid')
@@ -907,8 +906,7 @@ def axvline(self, x=0, ymin=0, ymax=1, **kwargs):
         xmin, xmax = self.get_xbound()
 
         # Strip away the units for comparison with non-unitized bounds.
-        self._process_unit_info({"x": x}, kwargs)
-        xx = self.convert_xunits(x)
+        xx, = self._process_unit_info({"x": x}, kwargs)
         scalex = (xx < xmin) or (xx > xmax)
 
         trans = self.get_xaxis_transform(which='grid')
@@ -1048,11 +1046,9 @@ def axhspan(self, ymin, ymax, xmin=0, xmax=1, **kwargs):
         self._check_no_units([xmin, xmax], ['xmin', 'xmax'])
         trans = self.get_yaxis_transform(which='grid')
 
-        self._process_unit_info({"x": [xmin, xmax], "y": [ymin, ymax]}, kwargs)
-
         # first we need to strip away the units
-        xmin, xmax = self.convert_xunits([xmin, xmax])
-        ymin, ymax = self.convert_yunits([ymin, ymax])
+        (xmin, xmax), (ymin, ymax) = self._process_unit_info(
+            {"x": [xmin, xmax], "y": [ymin, ymax]}, kwargs)
 
         verts = (xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin)
         p = mpatches.Polygon(verts, **kwargs)
@@ -1108,11 +1104,9 @@ def axvspan(self, xmin, xmax, ymin=0, ymax=1, **kwargs):
         self._check_no_units([ymin, ymax], ['ymin', 'ymax'])
         trans = self.get_xaxis_transform(which='grid')
 
-        self._process_unit_info({"x": [xmin, xmax], "y": [ymin, ymax]}, kwargs)
-
         # first we need to strip away the units
-        xmin, xmax = self.convert_xunits([xmin, xmax])
-        ymin, ymax = self.convert_yunits([ymin, ymax])
+        (xmin, xmax), (ymin, ymax) = self._process_unit_info(
+            {"x": [xmin, xmax], "y": [ymin, ymax]}, kwargs)
 
         verts = [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin)]
         p = mpatches.Polygon(verts, **kwargs)
@@ -1158,10 +1152,8 @@ def hlines(self, y, xmin, xmax, colors=None, linestyles='solid',
         """
 
         # We do the conversion first since not all unitized data is uniform
-        self._process_unit_info({"x": [xmin, xmax], "y": y}, kwargs)
-        y = self.convert_yunits(y)
-        xmin = self.convert_xunits(xmin)
-        xmax = self.convert_xunits(xmax)
+        (xmin, xmax), y = self._process_unit_info(
+            {"x": [xmin, xmax], "y": y}, kwargs)
 
         if not np.iterable(y):
             y = [y]
@@ -1239,10 +1231,8 @@ def vlines(self, x, ymin, ymax, colors=None, linestyles='solid',
         """
 
         # We do the conversion first since not all unitized data is uniform
-        self._process_unit_info({"x": x, "y": [ymin, ymax]}, kwargs)
-        x = self.convert_xunits(x)
-        ymin = self.convert_yunits(ymin)
-        ymax = self.convert_yunits(ymax)
+        x, (ymin, ymax) = self._process_unit_info(
+            {"x": x, "y": [ymin, ymax]}, kwargs)
 
         if not np.iterable(x):
             x = [x]
@@ -1380,13 +1370,9 @@ def eventplot(self, positions, orientation='horizontal', lineoffsets=1,
         --------
         .. plot:: gallery/lines_bars_and_markers/eventplot_demo.py
         """
-        self._process_unit_info(
-            {"x": positions, "y": [lineoffsets, linelengths]}, kwargs)
-
         # We do the conversion first since not all unitized data is uniform
-        positions = self.convert_xunits(positions)
-        lineoffsets = self.convert_yunits(lineoffsets)
-        linelengths = self.convert_yunits(linelengths)
+        positions, (lineoffsets, linelengths) = self._process_unit_info(
+            {"x": positions, "y": [lineoffsets, linelengths]}, kwargs)
 
         if not np.iterable(positions):
             positions = [positions]
@@ -2409,11 +2395,13 @@ def bar(self, x, height, width=0.8, bottom=None, *, align="center",
                 x = 0
 
         if orientation == 'vertical':
-            self._process_unit_info({"x": x, "y": height}, kwargs)
+            self._process_unit_info(
+                {"x": x, "y": height}, kwargs, convert=False)
             if log:
                 self.set_yscale('log', nonpositive='clip')
         elif orientation == 'horizontal':
-            self._process_unit_info({"x": width, "y": y}, kwargs)
+            self._process_unit_info(
+                {"x": width, "y": y}, kwargs, convert=False)
             if log:
                 self.set_xscale('log', nonpositive='clip')
 
@@ -2693,7 +2681,8 @@ def broken_barh(self, xranges, yrange, **kwargs):
             ydata = cbook.safe_first_element(yrange)
         else:
             ydata = None
-        self._process_unit_info({"x": xdata, "y": ydata}, kwargs)
+        self._process_unit_info(
+            {"x": xdata, "y": ydata}, kwargs, convert=False)
         xranges_conv = []
         for xr in xranges:
             if len(xr) != 2:
@@ -2814,13 +2803,9 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0,
             locs, heads, *args = args
 
         if orientation == 'vertical':
-            self._process_unit_info({"x": locs, "y": heads})
-            locs = self.convert_xunits(locs)
-            heads = self.convert_yunits(heads)
+            locs, heads = self._process_unit_info({"x": locs, "y": heads})
         else:
-            self._process_unit_info({"x": heads, "y": locs})
-            heads = self.convert_xunits(heads)
-            locs = self.convert_yunits(locs)
+            heads, locs = self._process_unit_info({"x": heads, "y": locs})
 
         # defaults for formats
         if linefmt is None:
@@ -3304,7 +3289,7 @@ def errorbar(self, x, y, yerr=None, xerr=None,
         if int(offset) != offset:
             raise ValueError("errorevery's starting index must be an integer")
 
-        self._process_unit_info({"x": x, "y": y}, kwargs)
+        self._process_unit_info({"x": x, "y": y}, kwargs, convert=False)
 
         # Make sure all the args are iterable; use lists not arrays to preserve
         # units.
@@ -4472,9 +4457,7 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
         """
         # Process **kwargs to handle aliases, conflicts with explicit kwargs:
 
-        self._process_unit_info({"x": x, "y": y}, kwargs)
-        x = self.convert_xunits(x)
-        y = self.convert_yunits(y)
+        x, y = self._process_unit_info({"x": x, "y": y}, kwargs)
 
         # np.ma.ravel yields an ndarray, not a masked array,
         # unless its argument is a masked array.
@@ -4703,7 +4686,7 @@ def reduce_C_function(C: array) -> float
             %(PolyCollection)s
 
         """
-        self._process_unit_info({"x": x, "y": y}, kwargs)
+        self._process_unit_info({"x": x, "y": y}, kwargs, convert=False)
 
         x, y, C = cbook.delete_masked_points(x, y, C)
 
@@ -5052,9 +5035,7 @@ def quiverkey(self, Q, X, Y, U, label, **kw):
     def _quiver_units(self, args, kw):
         if len(args) > 3:
             x, y = args[0:2]
-            self._process_unit_info({"x": x, "y": y}, kw)
-            x = self.convert_xunits(x)
-            y = self.convert_yunits(y)
+            x, y = self._process_unit_info({"x": x, "y": y}, kw)
             return (x, y) + args[2:]
         return args
 
@@ -5240,8 +5221,9 @@ def _fill_between_x_or_y(
                     self._get_patches_for_fill.get_next_color()
 
         # Handle united data, such as dates
-        self._process_unit_info({ind_dir: ind, dep_dir: dep1}, kwargs)
-        self._process_unit_info({dep_dir: dep2})
+        self._process_unit_info({ind_dir: ind, dep_dir: dep1}, kwargs,
+                                convert=False)
+        self._process_unit_info({dep_dir: dep2}, convert=False)
 
         # Convert the arrays so we can work with them
         ind = ma.masked_invalid(getattr(self, f"convert_{ind_dir}units")(ind))
@@ -5855,9 +5837,7 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
         Ny, Nx = X.shape
 
         # unit conversion allows e.g. datetime objects as axis values
-        self._process_unit_info({"x": X, "y": Y}, kwargs)
-        X = self.convert_xunits(X)
-        Y = self.convert_yunits(Y)
+        X, Y = self._process_unit_info({"x": X, "y": Y}, kwargs)
 
         # convert to MA, if necessary.
         C = ma.asarray(C)
@@ -6132,9 +6112,7 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         X = X.ravel()
         Y = Y.ravel()
         # unit conversion allows e.g. datetime objects as axis values
-        self._process_unit_info({"x": X, "y": Y}, kwargs)
-        X = self.convert_xunits(X)
-        Y = self.convert_yunits(Y)
+        X, Y = self._process_unit_info({"x": X, "y": Y}, kwargs)
 
         # convert to one dimensional arrays
         C = C.ravel()
@@ -6617,7 +6595,8 @@ def hist(self, x, bins=None, range=None, density=False, weights=None,
         # Process unit information
         # Unit conversion is done individually on each dataset
         self._process_unit_info(
-            {"x" if orientation == "vertical" else "y": x[0]}, kwargs)
+            {"x" if orientation == "vertical" else "y": x[0]}, kwargs,
+            convert=False)
         x = [self.convert_xunits(xi) for xi in x]
 
         if bin_range is not None:
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2175,9 +2175,10 @@ def update_datalim_bounds(self, bounds):
         """
         self.dataLim.set(mtransforms.Bbox.union([self.dataLim, bounds]))
 
-    def _process_unit_info(self, datasets=None, kwargs=None):
+    def _process_unit_info(self, datasets=None, kwargs=None, *, convert=True):
         """
-        Set axis units based on *datasets* and *kwargs*.
+        Set axis units based on *datasets* and *kwargs*, and optionally apply
+        unit conversions to *datasets*.
 
         Parameters
         ----------
@@ -2186,8 +2187,18 @@ def _process_unit_info(self, datasets=None, kwargs=None):
         kwargs : dict
             Other parameters from which unit info may be popped.  Note that
             this dict is mutated in-place!
+            Other parameters from which unit info may be popped.
+        convert : bool; default: True
+            Whether to return the original datasets or the converted ones.
+
+        Returns
+        -------
+        list
+            Either the original datasets (``datasets.values()``) if *convert*
+            is False, or the converted ones if *convert* is True (the default).
         """
         datasets = datasets or {}
+        retval = []
         axis_map = self._get_axis_map()
         for axis_name in datasets:
             if axis_name not in axis_map:
@@ -2213,7 +2224,9 @@ def _process_unit_info(self, datasets=None, kwargs=None):
                     # we need to update.
                     if data is not None:
                         axis.update_units(data)
-        return kwargs
+            if data is not None:
+                retval.append(axis.convert_units(data) if convert else data)
+        return retval
 
     def in_axes(self, mouseevent):
         """
@@ -3311,7 +3324,7 @@ def set_xlim(self, left=None, right=None, emit=True, auto=False,
                 raise TypeError('Cannot pass both `xmax` and `right`')
             right = xmax
 
-        self._process_unit_info({"x": (left, right)})
+        self._process_unit_info({"x": (left, right)}, convert=False)
         left = self._validate_converted_limits(left, self.convert_xunits)
         right = self._validate_converted_limits(right, self.convert_xunits)
 
@@ -3582,7 +3595,7 @@ def set_ylim(self, bottom=None, top=None, emit=True, auto=False,
                 raise TypeError('Cannot pass both `ymax` and `top`')
             top = ymax
 
-        self._process_unit_info({"y": (bottom, top)})
+        self._process_unit_info({"y": (bottom, top)}, convert=False)
         bottom = self._validate_converted_limits(bottom, self.convert_yunits)
         top = self._validate_converted_limits(top, self.convert_yunits)
 
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -1474,9 +1474,8 @@ def _check_xyz(self, args, kwargs):
         convert them to 2D using meshgrid.
         """
         x, y = args[:2]
-        self.axes._process_unit_info({"x": x, "y": y}, kwargs)
-        x = self.axes.convert_xunits(x)
-        y = self.axes.convert_yunits(y)
+        x, y = self.axes._process_unit_info(
+            {"x": x, "y": y}, kwargs, convert=True)
 
         x = np.asarray(x, dtype=np.float64)
         y = np.asarray(y, dtype=np.float64)
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -723,7 +723,7 @@ def set_xlim3d(self, left=None, right=None, emit=True, auto=False,
                 raise TypeError('Cannot pass both `xmax` and `right`')
             right = xmax
 
-        self._process_unit_info({"x": (left, right)})
+        self._process_unit_info({"x": (left, right)}, convert=False)
         left = self._validate_converted_limits(left, self.convert_xunits)
         right = self._validate_converted_limits(right, self.convert_xunits)
 
@@ -777,7 +777,7 @@ def set_ylim3d(self, bottom=None, top=None, emit=True, auto=False,
                 raise TypeError('Cannot pass both `ymax` and `top`')
             top = ymax
 
-        self._process_unit_info({"y": (bottom, top)})
+        self._process_unit_info({"y": (bottom, top)}, convert=False)
         bottom = self._validate_converted_limits(bottom, self.convert_yunits)
         top = self._validate_converted_limits(top, self.convert_yunits)
 
@@ -832,7 +832,7 @@ def set_zlim3d(self, bottom=None, top=None, emit=True, auto=False,
                 raise TypeError('Cannot pass both `zmax` and `top`')
             top = zmax
 
-        self._process_unit_info({"z": (bottom, top)})
+        self._process_unit_info({"z": (bottom, top)}, convert=False)
         bottom = self._validate_converted_limits(bottom, self.convert_zunits)
         top = self._validate_converted_limits(top, self.convert_zunits)
 
