Add helper function

dstansby · dstansby · commit 3fdb8d7a7d55 · 2018-04-13T16:39:36.000+01:00
diff --git a/lib/matplotlib/units.py b/lib/matplotlib/units.py
@@ -75,34 +75,32 @@ def wrapper(*args, **kwargs):
             if has_data:
                 data = arguments['data']
 
+            # Helper method to process unit info, and convert *original_data*
+            def _process_info(original_data, axis):
+                if axis == 'x':
+                    axes._process_unit_info(xdata=original_data, kwargs=kwargs)
+                    converted_data = axes.convert_xunits(original_data)
+                elif axis == 'y':
+                    axes._process_unit_info(ydata=original_data, kwargs=kwargs)
+                    converted_data = axes.convert_yunits(original_data)
+                return converted_data
+
             # Loop through each argument to be converted, update the axis
             # unit info, convert argument, and replace in *arguments* with
             # converted values
             for arg in convert_x:
                 if has_data and arguments[arg] in data:
                     data_arg = arguments[arg]
-                    original_data = data[data_arg]
-                    axes._process_unit_info(xdata=original_data, kwargs=kwargs)
-                    converted_data = axes.convert_xunits(original_data)
-                    data[data_arg] = converted_data
+                    data[data_arg] = _process_info(data[data_arg], 'x')
                 else:
-                    original_data = arguments[arg]
-                    axes._process_unit_info(xdata=original_data, kwargs=kwargs)
-                    converted_data = axes.convert_xunits(original_data)
-                    arguments[arg] = converted_data
+                    arguments[arg] = _process_info(arguments[arg] 'x')
 
             for arg in convert_y:
                 if has_data and arguments[arg] in data:
                     data_arg = arguments[arg]
-                    original_data = data[data_arg]
-                    axes._process_unit_info(ydata=original_data, kwargs=kwargs)
-                    converted_data = axes.convert_yunits(original_data)
-                    data[data_arg] = converted_data
+                    data[data_arg] = _process_info(data[data_arg], 'y')
                 else:
-                    original_data = arguments[arg]
-                    axes._process_unit_info(ydata=original_data, kwargs=kwargs)
-                    converted_data = axes.convert_yunits(original_data)
-                    arguments[arg] = converted_data
+                    arguments[arg] = _process_info(arguments[arg] 'x')
 
             if has_data:
                 arguments['data'] = data
