Factor out common code in _process_unit_info

dstansby · dstansby · commit 7e9d3053352d · 2018-03-04T13:58:02.000Z
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2117,40 +2117,32 @@ 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)
+        def _process_single_axis(data, axis, unit_name, kwargs):
+            if axis is None:
+                return
 
-        # 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)
+            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)
+
+            # Process kwargs 2nd since these will override default units
+            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):
