TST: Add test for pint (and alike) behavior (#8910)

This tests that unit libraries that wrap arrays, like pint, work properly. This adds an image test that checks current behavior, which seems to be fully correct.
matplotlib · dstansby · Aug 29, 2017 · Aug 17, 2017 · Aug 17, 2017 · Aug 18, 2017
commit 6eb96c80d107579c45f43f6912b9ad2530869a9c
diff --git a/lib/matplotlib/tests/baseline_images/test_units/plot_pint.png b/lib/matplotlib/tests/baseline_images/test_units/plot_pint.png
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -1,3 +1,4 @@
+from matplotlib.cbook import iterable
 import matplotlib.pyplot as plt
 from matplotlib.testing.decorators import image_comparison
 import matplotlib.units as munits
@@ -17,42 +18,65 @@ def __init__(self, data, units):
         self.units = units
 
     def to(self, new_units):
-        return Quantity(self.magnitude, new_units)
+        factors = {('hours', 'seconds'): 3600, ('minutes', 'hours'): 1 / 60,
+                   ('minutes', 'seconds'): 60, ('feet', 'miles'): 1 / 5280.,
+                   ('feet', 'inches'): 12, ('miles', 'inches'): 12 * 5280}
+        if self.units != new_units:
+            mult = factors[self.units, new_units]
+            return Quantity(mult * self.magnitude, new_units)
+        else:
+            return Quantity(self.magnitude, self.units)
 
     def __getattr__(self, attr):
         return getattr(self.magnitude, attr)
 
     def __getitem__(self, item):
-        return self.magnitude[item]
+        return Quantity(self.magnitude[item], self.units)
+
+    def __array__(self):
+        return np.asarray(self.magnitude)
 
 
 # Tests that the conversion machinery works properly for classes that
 # work as a facade over numpy arrays (like pint)
+@image_comparison(baseline_images=['plot_pint'],
+                  extensions=['png'], remove_text=False, style='mpl20')
 def test_numpy_facade():
     # 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)
+            return value.to(unit).magnitude
+        elif iterable(value):
+            try:
+                return [v.to(unit).magnitude for v in value]
+            except AttributeError:
+                return [Quantity(v, axis.get_units()).to(unit).magnitude
+                        for v in value]
         else:
             return Quantity(value, axis.get_units()).to(unit).magnitude
 
     qc.convert = MagicMock(side_effect=convert)
-    qc.axisinfo = MagicMock(return_value=None)
+    qc.axisinfo = MagicMock(side_effect=lambda u, a: munits.AxisInfo(label=u))
     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')
+    y = Quantity(np.linspace(0, 30), 'miles')
+    x = Quantity(np.linspace(0, 5), 'hours')
 
-    fig, ax = plt.subplots(1, 1)
-    l, = plt.plot(t, d)
-    ax.yaxis.set_units('inch')
+    fig, ax = plt.subplots()
+    fig.subplots_adjust(left=0.15)  # Make space for label
+    ax.plot(x, y, 'tab:blue')
+    ax.axhline(Quantity(26400, 'feet'), color='tab:red')
+    ax.axvline(Quantity(120, 'minutes'), color='tab:green')
+    ax.yaxis.set_units('inches')
+    ax.xaxis.set_units('seconds')
+    ax.autoscale_view()
 
     assert qc.convert.called
     assert qc.axisinfo.called