Proof of concept "accepts units" decorator

dstansby · dstansby · commit cffa3643ff24 · 2018-02-14T22:20:00.000Z
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -4059,7 +4059,7 @@ def dopatch(xs, ys, **kwargs):
         return dict(whiskers=whiskers, caps=caps, boxes=boxes,
                     medians=medians, fliers=fliers, means=means)
 
-    @munits._accepts_units(self, convert_x=[1], convert_y=[2])
+    @munits._accepts_units(convert_x=['x'], convert_y=['y'])
     @_preprocess_data(replace_names=["x", "y", "s", "linewidths",
                                      "edgecolors", "c", "facecolor",
                                      "facecolors", "color"],
@@ -4211,7 +4211,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'
 
-
         # 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
@@ -44,12 +44,81 @@ def default_units(x, axis):
 from __future__ import (absolute_import, division, print_function,
                         unicode_literals)
 
-
+import inspect
 import six
 from matplotlib.cbook import iterable, is_numlike, safe_first_element
 import numpy as np
 
 
+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']
+
+            # 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
+                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
+
+            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
+                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
+
+            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
