|
| 1 | +from matplotlib.cbook import iterable |
1 | 2 | import matplotlib.pyplot as plt |
| 3 | +from matplotlib.testing.decorators import image_comparison |
2 | 4 | import matplotlib.units as munits |
3 | 5 | import numpy as np |
4 | 6 |
|
|
9 | 11 | from mock import MagicMock |
10 | 12 |
|
11 | 13 |
|
12 | | -# Tests that the conversion machinery works properly for classes that |
13 | | -# work as a facade over numpy arrays (like pint) |
14 | | -def test_numpy_facade(): |
15 | | - # Basic class that wraps numpy array and has units |
16 | | - class Quantity(object): |
17 | | - def __init__(self, data, units): |
18 | | - self.magnitude = data |
19 | | - self.units = units |
| 14 | +# Basic class that wraps numpy array and has units |
| 15 | +class Quantity(object): |
| 16 | + def __init__(self, data, units): |
| 17 | + self.magnitude = data |
| 18 | + self.units = units |
| 19 | + |
| 20 | + def to(self, new_units): |
| 21 | + factors = {('hours', 'seconds'): 3600, ('minutes', 'hours'): 1 / 60, |
| 22 | + ('minutes', 'seconds'): 60, ('feet', 'miles'): 1 / 5280., |
| 23 | + ('feet', 'inches'): 12, ('miles', 'inches'): 12 * 5280} |
| 24 | + if self.units != new_units: |
| 25 | + mult = factors[self.units, new_units] |
| 26 | + return Quantity(mult * self.magnitude, new_units) |
| 27 | + else: |
| 28 | + return Quantity(self.magnitude, self.units) |
| 29 | + |
| 30 | + def __getattr__(self, attr): |
| 31 | + return getattr(self.magnitude, attr) |
20 | 32 |
|
21 | | - def to(self, new_units): |
22 | | - return Quantity(self.magnitude, new_units) |
| 33 | + def __getitem__(self, item): |
| 34 | + return Quantity(self.magnitude[item], self.units) |
23 | 35 |
|
24 | | - def __getattr__(self, attr): |
25 | | - return getattr(self.magnitude, attr) |
| 36 | + def __array__(self): |
| 37 | + return np.asarray(self.magnitude) |
26 | 38 |
|
27 | | - def __getitem__(self, item): |
28 | | - return self.magnitude[item] |
29 | 39 |
|
| 40 | +# Tests that the conversion machinery works properly for classes that |
| 41 | +# work as a facade over numpy arrays (like pint) |
| 42 | +@image_comparison(baseline_images=['plot_pint'], |
| 43 | + extensions=['png'], remove_text=False, style='mpl20') |
| 44 | +def test_numpy_facade(): |
30 | 45 | # Create an instance of the conversion interface and |
31 | 46 | # mock so we can check methods called |
32 | 47 | qc = munits.ConversionInterface() |
33 | 48 |
|
34 | 49 | def convert(value, unit, axis): |
35 | 50 | if hasattr(value, 'units'): |
36 | | - return value.to(unit) |
| 51 | + return value.to(unit).magnitude |
| 52 | + elif iterable(value): |
| 53 | + try: |
| 54 | + return [v.to(unit).magnitude for v in value] |
| 55 | + except AttributeError: |
| 56 | + return [Quantity(v, axis.get_units()).to(unit).magnitude |
| 57 | + for v in value] |
37 | 58 | else: |
38 | 59 | return Quantity(value, axis.get_units()).to(unit).magnitude |
39 | 60 |
|
40 | 61 | qc.convert = MagicMock(side_effect=convert) |
41 | | - qc.axisinfo = MagicMock(return_value=None) |
| 62 | + qc.axisinfo = MagicMock(side_effect=lambda u, a: munits.AxisInfo(label=u)) |
42 | 63 | qc.default_units = MagicMock(side_effect=lambda x, a: x.units) |
43 | 64 |
|
44 | 65 | # Register the class |
45 | 66 | munits.registry[Quantity] = qc |
46 | 67 |
|
47 | 68 | # Simple test |
48 | | - t = Quantity(np.linspace(0, 10), 'sec') |
49 | | - d = Quantity(30 * np.linspace(0, 10), 'm/s') |
| 69 | + y = Quantity(np.linspace(0, 30), 'miles') |
| 70 | + x = Quantity(np.linspace(0, 5), 'hours') |
50 | 71 |
|
51 | | - fig, ax = plt.subplots(1, 1) |
52 | | - l, = plt.plot(t, d) |
53 | | - ax.yaxis.set_units('inch') |
| 72 | + fig, ax = plt.subplots() |
| 73 | + fig.subplots_adjust(left=0.15) # Make space for label |
| 74 | + ax.plot(x, y, 'tab:blue') |
| 75 | + ax.axhline(Quantity(26400, 'feet'), color='tab:red') |
| 76 | + ax.axvline(Quantity(120, 'minutes'), color='tab:green') |
| 77 | + ax.yaxis.set_units('inches') |
| 78 | + ax.xaxis.set_units('seconds') |
54 | 79 |
|
55 | 80 | assert qc.convert.called |
56 | 81 | assert qc.axisinfo.called |
57 | 82 | assert qc.default_units.called |
| 83 | + |
| 84 | + |
| 85 | +# Tests gh-8908 |
| 86 | +@image_comparison(baseline_images=['plot_masked_units'], |
| 87 | + extensions=['png'], remove_text=True, style='mpl20') |
| 88 | +def test_plot_masked_units(): |
| 89 | + data = np.linspace(-5, 5) |
| 90 | + data_masked = np.ma.array(data, mask=(data > -2) & (data < 2)) |
| 91 | + data_masked_units = Quantity(data_masked, 'meters') |
| 92 | + |
| 93 | + fig, ax = plt.subplots() |
| 94 | + ax.plot(data_masked_units) |
0 commit comments