List-based _process_unit_info API.

anntzer · anntzer · commit bb5e68b727f5 · 2020-09-21T02:03:36.000+02:00
Better if multiple datasets need to be converted at once...
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -709,7 +709,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.
-        yy, = self._process_unit_info({"y": y}, kwargs)
+        yy, = self._process_unit_info([("y", y)], kwargs)
         scaley = (yy < ymin) or (yy > ymax)
 
         trans = self.get_yaxis_transform(which='grid')
@@ -776,7 +776,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.
-        xx, = self._process_unit_info({"x": x}, kwargs)
+        xx, = self._process_unit_info([("x", x)], kwargs)
         scalex = (xx < xmin) or (xx > xmax)
 
         trans = self.get_xaxis_transform(which='grid')
@@ -918,7 +918,7 @@ def axhspan(self, ymin, ymax, xmin=0, xmax=1, **kwargs):
 
         # first we need to strip away the units
         (xmin, xmax), (ymin, ymax) = self._process_unit_info(
-            {"x": [xmin, xmax], "y": [ymin, ymax]}, kwargs)
+            [("x", [xmin, xmax]), ("y", [ymin, ymax])], kwargs)
 
         verts = (xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin)
         p = mpatches.Polygon(verts, **kwargs)
@@ -976,7 +976,7 @@ def axvspan(self, xmin, xmax, ymin=0, ymax=1, **kwargs):
 
         # first we need to strip away the units
         (xmin, xmax), (ymin, ymax) = self._process_unit_info(
-            {"x": [xmin, xmax], "y": [ymin, ymax]}, kwargs)
+            [("x", [xmin, xmax]), ("y", [ymin, ymax])], kwargs)
 
         verts = [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin)]
         p = mpatches.Polygon(verts, **kwargs)
@@ -1022,11 +1022,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, convert=False)
-        y = self.convert_yunits(y)
-        xmin = self.convert_xunits(xmin)
-        xmax = self.convert_xunits(xmax)
+        xmin, xmax, y = self._process_unit_info(
+            [("x", xmin), ("x", xmax), ("y", y)], kwargs)
 
         if not np.iterable(y):
             y = [y]
@@ -1102,11 +1099,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, convert=False)
-        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), ("y", ymax)], kwargs)
 
         if not np.iterable(x):
             x = [x]
@@ -1245,8 +1239,8 @@ def eventplot(self, positions, orientation='horizontal', lineoffsets=1,
         .. plot:: gallery/lines_bars_and_markers/eventplot_demo.py
         """
         # We do the conversion first since not all unitized data is uniform
-        positions, (lineoffsets, linelengths) = self._process_unit_info(
-            {"x": positions, "y": [lineoffsets, linelengths]}, kwargs)
+        positions, lineoffsets, linelengths = self._process_unit_info(
+            [("x", positions), ("y", lineoffsets), ("y", linelengths)], kwargs)
 
         if not np.iterable(positions):
             positions = [positions]
@@ -2269,12 +2263,12 @@ def bar(self, x, height, width=0.8, bottom=None, *, align="center",
 
         if orientation == 'vertical':
             self._process_unit_info(
-                {"x": x, "y": height}, kwargs, convert=False)
+                [("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, convert=False)
+                [("x", width), ("y", y)], kwargs, convert=False)
             if log:
                 self.set_xscale('log', nonpositive='clip')
 
@@ -2555,7 +2549,7 @@ def broken_barh(self, xranges, yrange, **kwargs):
         else:
             ydata = None
         self._process_unit_info(
-            {"x": xdata, "y": ydata}, kwargs, convert=False)
+            [("x", xdata), ("y", ydata)], kwargs, convert=False)
         xranges_conv = []
         for xr in xranges:
             if len(xr) != 2:
@@ -2675,9 +2669,9 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0,
             locs, heads, *args = args
 
         if orientation == 'vertical':
-            locs, heads = self._process_unit_info({"x": locs, "y": heads})
+            locs, heads = self._process_unit_info([("x", locs), ("y", heads)])
         else:
-            heads, locs = self._process_unit_info({"x": heads, "y": locs})
+            heads, locs = self._process_unit_info([("x", heads), ("y", locs)])
 
         # defaults for formats
         if linefmt is None:
@@ -3161,7 +3155,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, convert=False)
+        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.
@@ -4328,7 +4322,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:
 
-        x, y = self._process_unit_info({"x": x, "y": y}, kwargs)
+        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.
@@ -4557,7 +4551,7 @@ def reduce_C_function(C: array) -> float
             %(PolyCollection)s
 
         """
-        self._process_unit_info({"x": x, "y": y}, kwargs, convert=False)
+        self._process_unit_info([("x", x), ("y", y)], kwargs, convert=False)
 
         x, y, C = cbook.delete_masked_points(x, y, C)
 
@@ -4906,7 +4900,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]
-            x, y = self._process_unit_info({"x": x, "y": y}, kw)
+            x, y = self._process_unit_info([("x", x), ("y", y)], kw)
             return (x, y) + args[2:]
         return args
 
@@ -5092,9 +5086,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,
-                                convert=False)
-        self._process_unit_info({dep_dir: dep2}, convert=False)
+        self._process_unit_info(
+            [(ind_dir, ind), (dep_dir, dep1), (dep_dir, dep2)], kwargs,
+            convert=False)
 
         # Convert the arrays so we can work with them
         ind = ma.masked_invalid(getattr(self, f"convert_{ind_dir}units")(ind))
@@ -5716,7 +5710,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
-        X, Y = self._process_unit_info({"x": X, "y": Y}, kwargs)
+        X, Y = self._process_unit_info([("x", X), ("y", Y)], kwargs)
 
         # convert to MA, if necessary.
         C = ma.asarray(C)
@@ -5991,7 +5985,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
-        X, Y = self._process_unit_info({"x": X, "y": Y}, kwargs)
+        X, Y = self._process_unit_info([("x", X), ("y", Y)], kwargs)
 
         # convert to one dimensional arrays
         C = C.ravel()
@@ -6475,11 +6469,11 @@ def hist(self, x, bins=None, range=None, density=False, weights=None,
         # one at a time.
         if orientation == "vertical":
             convert_units = self.convert_yunits
-            x = [*self._process_unit_info({"x": x[0]}, kwargs),
+            x = [*self._process_unit_info([("x", x[0])], kwargs),
                  *map(convert_units, x[1:])]
         else:  # horizontal
             convert_units = self.convert_yunits
-            x = [*self._process_unit_info({"y": x[0]}, kwargs),
+            x = [*self._process_unit_info([("y", x[0])], kwargs),
                  *map(convert_units, x[1:])]
 
         if bin_range is not None:
@@ -6768,7 +6762,7 @@ def stairs(self, values, edges=None, *,
             edges = np.arange(len(values) + 1)
 
         edges, values = self._process_unit_info(
-            {"x": edges, "y": values}, kwargs)
+            [("x", edges), ("y", values)], kwargs)
 
         patch = mpatches.StepPatch(values,
                                    edges,
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2215,8 +2215,9 @@ def _process_unit_info(self, datasets=None, kwargs=None, *, convert=True):
 
         Parameters
         ----------
-        datasets : dict
-            Mapping of axis names (per `._get_axis_map`) to datasets.
+        datasets : list
+            List of (axis_name, dataset) pairs (where the axis name is defined
+            as in `._get_axis_map`.
         kwargs : dict
             Other parameters from which unit info (i.e., the *xunits*,
             *yunits*, *zunits* (for 3D axes), *runits* and *thetaunits* (for
@@ -2228,39 +2229,43 @@ def _process_unit_info(self, datasets=None, kwargs=None, *, convert=True):
         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 = []
+            Either the original datasets if *convert* is False, or the
+            converted ones if *convert* is True (the default).
+        """
+        # The API makes datasets a list of pairs rather than an axis_name to
+        # dataset mapping because it is sometimes necessary to process multiple
+        # datasets for a single axis, and concatenating them may be tricky
+        # (e.g. if some are scalars, etc.).
+        datasets = datasets or []
+        kwargs = kwargs or {}
         axis_map = self._get_axis_map()
-        for axis_name in datasets:
-            if axis_name not in axis_map:
-                raise ValueError(f"Invalid axis name: {axis_name!r}")
+        for axis_name, data in datasets:
+            try:
+                axis = axis_map[axis_name]
+            except KeyError:
+                raise ValueError(f"Invalid axis name: {axis_name!r}") from None
+            # Update from data if axis is already set but no unit is set yet.
+            if axis is not None and data is not None and not axis.have_units():
+                axis.update_units(data)
         for axis_name, axis in axis_map.items():
             # Return if no axis is set.
             if axis is None:
                 continue
-            data = datasets.get(axis_name)
-            # Update from data if no unit is set yet.
-            if data is not None and not axis.have_units():
-                axis.update_units(data)
             # Check for units in the kwargs, and if present update axis.
-            if kwargs is not None:
-                units = kwargs.pop(f"{axis_name}units", axis.units)
-                if self.name == "polar":
-                    # Special case: polar supports "thetaunits"/"runits".
-                    polar_units = {"x": "thetaunits", "y": "runits"}
-                    units = kwargs.pop(polar_units[axis_name], units)
-                if units != axis.units and units is not None:
-                    axis.set_units(units)
-                    # If the units being set imply a different converter,
-                    # we need to update.
-                    if data is not None:
+            units = kwargs.pop(f"{axis_name}units", axis.units)
+            if self.name == "polar":
+                # Special case: polar supports "thetaunits"/"runits".
+                polar_units = {"x": "thetaunits", "y": "runits"}
+                units = kwargs.pop(polar_units[axis_name], units)
+            if units != axis.units and units is not None:
+                axis.set_units(units)
+                # If the units being set imply a different converter,
+                # we need to update again.
+                for dataset_axis_name, data in datasets:
+                    if dataset_axis_name == axis_name and data is not None:
                         axis.update_units(data)
-            if data is not None:
-                retval.append(axis.convert_units(data) if convert else data)
-        return retval
+        return [axis_map[axis_name].convert_units(data) if convert else data
+                for axis_name, data in datasets]
 
     def in_axes(self, mouseevent):
         """
@@ -3412,7 +3417,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)}, convert=False)
+        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)
 
@@ -3737,7 +3742,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)}, convert=False)
+        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
@@ -1476,7 +1476,7 @@ def _check_xyz(self, args, kwargs):
         """
         x, y = args[:2]
         x, y = self.axes._process_unit_info(
-            {"x": x, "y": y}, kwargs, convert=True)
+            [("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
@@ -724,7 +724,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)}, convert=False)
+        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)
 
@@ -778,7 +778,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)}, convert=False)
+        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)
 
@@ -833,7 +833,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)}, convert=False)
+        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)
 
