Add test for unit conversion.

dopplershift · dopplershift · commit 0925e38fc9fb · 2015-09-14T12:55:59.000-06:00
This is catching the problems with pint-like libraries who wrap numpy
arrays.
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
