Proof of concept "accepts units" decorator

dstansby · dstansby · commit 79f4a111a9be · 2018-02-10T19:00:12.000Z
diff --git a/.travis.yml b/.travis.yml
@@ -65,19 +65,6 @@ env:
 
 matrix:
   include:
-    - python: 2.7
-      # pytest-cov>=2.3.1 due to https://github.com/pytest-dev/pytest-cov/issues/124.
-      env:
-        - CYCLER=cycler==0.10
-        - DATEUTIL=python-dateutil==2.1
-        - MOCK=mock
-        - NOSE=nose
-        - NUMPY=numpy==1.7.1
-        - PANDAS='pandas<0.21.0'
-        - PYPARSING=pyparsing==2.0.1
-        - PYTEST=pytest==3.1.0
-        - PYTEST_COV=pytest-cov==2.3.1
-        - SPHINX=sphinx==1.3
     - python: 3.4
       env: PYTHON_ARGS=-OO
     - python: 3.6
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,70 @@ 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 (can be mixed). Integers denote
+        non-keyword arguments to be converted, and strings denote keyword
+        arguments to be converted.
+    """
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            axes = args[0]
+            bound_args = inspect.signature(func).bind(*args, **kwargs)
+            arguments = bound_args.arguments
+            has_data = (('data' in arguments) and
+                        (arguments['data'] is not None))
+            if has_data:
+                data = arguments['data']
+
+            for arg in convert_x:
+                if has_data and arg in data:
+                    original_data = data[arg]
+                    axes._process_unit_info(xdata=original_data, kwargs=kwargs)
+                    converted_data = axes.convert_xunits(original_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 arg in data:
+                    original_data = data[arg]
+                    axes._process_unit_info(ydata=original_data, kwargs=kwargs)
+                    converted_data = axes.convert_yunits(original_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
+            bound_args.arguments = arguments
+            args = bound_args.args
+            kwargs = bound_args.kwargs
+            return func(*args, **kwargs)
+        return wrapper
+    return decorator
+
+
 class AxisInfo(object):
     """
     Information to support default axis labeling, tick labeling, and
