Skip to content

Commit 6eb96c8

Browse files
committed
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.
1 parent a90064e commit 6eb96c8

File tree

2 files changed

+33
-9
lines changed

2 files changed

+33
-9
lines changed

lib/matplotlib/tests/test_units.py

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from matplotlib.cbook import iterable
12
import matplotlib.pyplot as plt
23
from matplotlib.testing.decorators import image_comparison
34
import matplotlib.units as munits
@@ -17,42 +18,65 @@ def __init__(self, data, units):
1718
self.units = units
1819

1920
def to(self, new_units):
20-
return Quantity(self.magnitude, 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)
2129

2230
def __getattr__(self, attr):
2331
return getattr(self.magnitude, attr)
2432

2533
def __getitem__(self, item):
26-
return self.magnitude[item]
34+
return Quantity(self.magnitude[item], self.units)
35+
36+
def __array__(self):
37+
return np.asarray(self.magnitude)
2738

2839

2940
# Tests that the conversion machinery works properly for classes that
3041
# work as a facade over numpy arrays (like pint)
42+
@image_comparison(baseline_images=['plot_pint'],
43+
extensions=['png'], remove_text=False, style='mpl20')
3144
def test_numpy_facade():
3245
# Create an instance of the conversion interface and
3346
# mock so we can check methods called
3447
qc = munits.ConversionInterface()
3548

3649
def convert(value, unit, axis):
3750
if hasattr(value, 'units'):
38-
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]
3958
else:
4059
return Quantity(value, axis.get_units()).to(unit).magnitude
4160

4261
qc.convert = MagicMock(side_effect=convert)
43-
qc.axisinfo = MagicMock(return_value=None)
62+
qc.axisinfo = MagicMock(side_effect=lambda u, a: munits.AxisInfo(label=u))
4463
qc.default_units = MagicMock(side_effect=lambda x, a: x.units)
4564

4665
# Register the class
4766
munits.registry[Quantity] = qc
4867

4968
# Simple test
50-
t = Quantity(np.linspace(0, 10), 'sec')
51-
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')
5271

53-
fig, ax = plt.subplots(1, 1)
54-
l, = plt.plot(t, d)
55-
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')
79+
ax.autoscale_view()
5680

5781
assert qc.convert.called
5882
assert qc.axisinfo.called

0 commit comments

Comments
 (0)