Merge pull request #4803 from dopplershift/fix-plot-units

tacaswell · tacaswell · commit c27970c97320 · 2015-09-14T15:28:31.000-04:00
FIX: unit support with `plot` and `pint`
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1465,6 +1465,7 @@ def tk_window_focus():
     'matplotlib.tests.test_tightlayout',
     'matplotlib.tests.test_transforms',
     'matplotlib.tests.test_triangulation',
+    'matplotlib.tests.test_units',
     'matplotlib.tests.test_widgets',
     'matplotlib.tests.test_cycles',
     'matplotlib.tests.test_labeled_data_unpacking',
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -15,7 +15,8 @@
 import matplotlib
 
 from matplotlib import cbook
-from matplotlib.cbook import _string_to_bool, iterable, index_of, get_label
+from matplotlib.cbook import (_check_1d, _string_to_bool, iterable,
+                              index_of, get_label)
 from matplotlib import docstring
 import matplotlib.colors as mcolors
 import matplotlib.lines as mlines
@@ -214,8 +215,10 @@ def _xy_from_xy(self, x, y):
                 if by:
                     y = self.axes.convert_yunits(y)
 
-        x = np.atleast_1d(x)  # like asanyarray, but converts scalar to array
-        y = np.atleast_1d(y)
+        # like asanyarray, but converts scalar to array, and doesn't change
+        # existing compatible sequences
+        x = _check_1d(x)
+        y = _check_1d(y)
         if x.shape[0] != y.shape[0]:
             raise ValueError("x and y must have same first dimension")
         if x.ndim > 2 or y.ndim > 2:
@@ -353,8 +356,8 @@ def _plot_args(self, tup, kwargs):
             kwargs['label'] = get_label(tup[-1], None)
 
         if len(tup) == 2:
-            x = np.atleast_1d(tup[0])
-            y = np.atleast_1d(tup[-1])
+            x = _check_1d(tup[0])
+            y = _check_1d(tup[-1])
         else:
             x, y = index_of(tup[-1])
 
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -2192,6 +2192,21 @@ def is_math_text(s):
     return even_dollars
 
 
+def _check_1d(x):
+    '''
+    Converts a sequence of less than 1 dimension, to an array of 1
+    dimension; leaves everything else untouched.
+    '''
+    if not hasattr(x, 'shape') or len(x.shape) < 1:
+        return np.atleast_1d(x)
+    else:
+        try:
+            x[:, None]
+            return x
+        except (IndexError, TypeError):
+            return np.atleast_1d(x)
+
+
 def _reshape_2D(X):
     """
     Converts a non-empty list or an ndarray of two or fewer dimensions
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -0,0 +1,57 @@
+import matplotlib.pyplot as plt
+import matplotlib.units as munits
+import numpy as np
+
+try:
+    # mock in python 3.3+
+    from unittest.mock import MagicMock
+except ImportError:
+    from mock import MagicMock
+
+
+# Tests that the conversion machinery works properly for classes that
+# work as a facade over numpy arrays (like pint)
+def test_numpy_facade():
+    # Basic class that wraps numpy array and has units
+    class Quantity(object):
+        def __init__(self, data, units):
+            self.magnitude = data
+            self.units = units
+
+        def to(self, new_units):
+            return Quantity(self.magnitude, new_units)
+
+        def __getattr__(self, attr):
+            return getattr(self.magnitude, attr)
+
+        def __getitem__(self, item):
+            return self.magnitude[item]
+
+    # Create an instance of the conversion interface and
+    # mock so we can check methods called
+    qc = munits.ConversionInterface()
+
+    def convert(value, unit, axis):
+        if hasattr(value, 'units'):
+            return value.to(unit)
+        else:
+            return Quantity(value, axis.get_units()).to(unit).magnitude
+
+    qc.convert = MagicMock(side_effect=convert)
+    qc.axisinfo = MagicMock(return_value=None)
+    qc.default_units = MagicMock(side_effect=lambda x, a: x.units)
+
+    # Register the class
+    munits.registry[Quantity] = qc
+
+    # Simple test
+    t = Quantity(np.linspace(0, 10), 'sec')
+    d = Quantity(30 * np.linspace(0, 10), 'm/s')
+
+    fig, ax = plt.subplots(1, 1)
+    l, = plt.plot(t, d)
+    ax.yaxis.set_units('inch')
+
+    assert qc.convert.called
+    assert qc.axisinfo.called
+    assert qc.default_units.called
