Proof of concept "accepts units" decorator

dstansby · dstansby · commit f44ee448d450 · 2018-04-13T16:55:46.000+01:00
Proof of concept "accepts units" decorator

Add helper function

Mssing commas
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -33,6 +33,7 @@
 import matplotlib.ticker as mticker
 import matplotlib.transforms as mtransforms
 import matplotlib.tri as mtri
+import matplotlib.units as munits
 from matplotlib.cbook import (
     mplDeprecation, warn_deprecated, STEP_LOOKUP_MAP, iterable,
     safe_first_element)
@@ -3844,6 +3845,7 @@ def dopatch(xs, ys, **kwargs):
         return dict(whiskers=whiskers, caps=caps, boxes=boxes,
                     medians=medians, fliers=fliers, means=means)
 
+    @munits._accepts_units(convert_x=['x'], convert_y=['y'])
     @_preprocess_data(replace_names=["x", "y", "s", "linewidths",
                                      "edgecolors", "c", "facecolor",
                                      "facecolors", "color"],
@@ -3990,10 +3992,6 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
         if edgecolors is None and not rcParams['_internal.classic_mode']:
             edgecolors = 'face'
 
-        self._process_unit_info(xdata=x, ydata=y, kwargs=kwargs)
-        x = self.convert_xunits(x)
-        y = self.convert_yunits(y)
-
         # np.ma.ravel yields an ndarray, not a masked array,
         # unless its argument is a masked array.
         xy_shape = (np.shape(x), np.shape(y))
diff --git a/lib/matplotlib/units.py b/lib/matplotlib/units.py
@@ -43,12 +43,80 @@ def default_units(x, axis):
 """
 
 from numbers import Number
+import inspect
 
 import numpy as np
 
 from matplotlib.cbook import iterable, safe_first_element
 
 
+def _accepts_units(convert_x, convert_y):
+    """
+    A decorator for functions and methods that accept units. The parameters
+    indicated in *convert_x* and *convert_y* are used to update the axis
+    unit information, are converted, and then handed on to the decorated
+    function.
+
+    Parameters
+    ----------
+    convert_x, convert_y : list
+        A list of integers or strings, indicating the arguments to be converted
+    """
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            axes = args[0]
+            # Bind the incoming arguments to the function signature
+            bound_args = inspect.signature(func).bind(*args, **kwargs)
+            # Get the original arguments - these will be modified later
+            arguments = bound_args.arguments
+            # Check for data kwarg
+            has_data = (('data' in arguments) and
+                        (arguments['data'] is not None))
+            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]
+                    data[data_arg] = _process_info(data[data_arg], 'x')
+                else:
+                    arguments[arg] = _process_info(arguments[arg], 'x')
+
+            for arg in convert_y:
+                if has_data and arguments[arg] in data:
+                    data_arg = arguments[arg]
+                    data[data_arg] = _process_info(data[data_arg], 'y')
+                else:
+                    arguments[arg] = _process_info(arguments[arg], 'y')
+
+            if has_data:
+                arguments['data'] = data
+            # Update the arguments with converted values
+            bound_args.arguments = arguments
+
+            # Give updated values to the original function
+            args = bound_args.args
+            kwargs = bound_args.kwargs
+            kwargs.pop('xunits', None)
+            kwargs.pop('yunits', None)
+            return func(*args, **kwargs)
+        return wrapper
+    return decorator
+
+
 class AxisInfo(object):
     """
     Information to support default axis labeling, tick labeling, and
