Merge pull request #10408 from dstansby/unit-common-code

timhoffm · web-flow · commit 92018b563207 · 2018-04-07T09:22:59.000+02:00
Factor out common code in _process_unit_info
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2060,40 +2060,33 @@ def update_datalim_bounds(self, bounds):
     def _process_unit_info(self, xdata=None, ydata=None, kwargs=None):
         """Look for unit *kwargs* and update the axis instances as necessary"""
 
-        if self.xaxis is None or self.yaxis is None:
-            return
-
-        if xdata is not None:
-            # we only need to update if there is nothing set yet.
-            if not self.xaxis.have_units():
-                self.xaxis.update_units(xdata)
-
-        if ydata is not None:
-            # we only need to update if there is nothing set yet.
-            if not self.yaxis.have_units():
-                self.yaxis.update_units(ydata)
-
-        # process kwargs 2nd since these will override default units
-        if kwargs is not None:
-            xunits = kwargs.pop('xunits', self.xaxis.units)
-            if self.name == 'polar':
-                xunits = kwargs.pop('thetaunits', xunits)
-            if xunits != self.xaxis.units:
-                self.xaxis.set_units(xunits)
-                # If the units being set imply a different converter,
-                # we need to update.
-                if xdata is not None:
-                    self.xaxis.update_units(xdata)
-
-            yunits = kwargs.pop('yunits', self.yaxis.units)
-            if self.name == 'polar':
-                yunits = kwargs.pop('runits', yunits)
-            if yunits != self.yaxis.units:
-                self.yaxis.set_units(yunits)
-                # If the units being set imply a different converter,
-                # we need to update.
-                if ydata is not None:
-                    self.yaxis.update_units(ydata)
+        def _process_single_axis(data, axis, unit_name, kwargs):
+            # Return if there's no axis set
+            if axis is None:
+                return kwargs
+
+            if data is not None:
+                # We only need to update if there is nothing set yet.
+                if 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(unit_name, axis.units)
+                if self.name == 'polar':
+                    polar_units = {'xunits': 'thetaunits', 'yunits': 'runits'}
+                    units = kwargs.pop(polar_units[unit_name], units)
+
+                if units != axis.units:
+                    axis.set_units(units)
+                    # If the units being set imply a different converter,
+                    # we need to update.
+                    if data is not None:
+                        axis.update_units(data)
+            return kwargs
+
+        kwargs = _process_single_axis(xdata, self.xaxis, 'xunits', kwargs)
+        kwargs = _process_single_axis(ydata, self.yaxis, 'yunits', kwargs)
         return kwargs
 
     def in_axes(self, mouseevent):
