10000 Check size of input array · matplotlib/matplotlib@0cfe854 · GitHub
[go: up one dir, main page]
Skip to content

Commit 0cfe854

Browse files
ericpreericpre
committed
Check size of input array
1 parent bdde2e9 commit 0cfe854

File tree

3 files changed

+68
-15
lines changed

3 files changed

+68
-15
lines changed

‎doc/users/next_whats_new/add_EllipseCollection_setters.rst

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,8 @@ can now be changed after the collection have been created.
3333
ax.set_xlim(-2, 12)
3434
ax.set_ylim(-2, 12)
3535

36-
new_widths = rng.random((10, 2)) * 2
37-
new_heights = rng.random((10, 2)) * 3
38-
new_angles = rng.random((10, 2)) * 180
36+
new_widths = rng.random((10, )) * 2
37+
new_heights = rng.random((10, )) * 3
38+
new_angles = rng.random((10, )) * 180
3939

4040
ec.set(widths=new_widths, heights=new_heights, angles=new_angles)

‎lib/matplotlib/collections.py

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1728,11 +1728,16 @@ def __init__(self, widths, heights, angles, *, units='points', **kwargs):
17281728
Parameters
17291729
----------
17301730
widths : array-like
1731-
The lengths of the first axes (e.g., major axis lengths).
1731+
The lengths of the first axes (e.g., major axis lengths). The size
1732+
of the array must one or the same as other array inputs
1733+
(``heights``, ``angles``).
17321734
heights : array-like
1733-
The lengths of second axes.
1735+
The lengths of second axes. The size of the array must one or
1736+
the same as other array inputs (``widths``, ``angles``).
17341737
angles : array-like
1735-
The angles of the first axes, degrees CCW from the x-axis.
1738+
The angles of the first axes, degrees CCW from the x-axis. The size
1739+
of the array must one or the same as other array inputs
1740+
(``widths``, ``angles``).
17361741
units : {'points', 'inches', 'dots', 'width', 'height', 'x', 'y', 'xy'}
17371742
The units in which majors and minors are given; 'width' and
17381743
'height' refer to the dimensions of the axes, while 'x' and 'y'
@@ -1745,14 +1750,25 @@ def __init__(self, widths, heights, angles, *, units='points', **kwargs):
17451750
Forwarded to `Collection`.
17461751
"""
17471752
super().__init__(**kwargs)
1753+
17481754
self._set_widths(widths)
17491755
self._set_heights(heights)
17501756
self._set_angles(angles)
17511757
self._units = units
1758+
17521759
self.set_transform(transforms.IdentityTransform())
17531760
self._transforms = np.empty((0, 3, 3))
17541761
self._paths = [mpath.Path.unit_circle()]
17551762

1763+
def _check_length(self, array, name):
1764+
N = len(self.get_offsets())
1765+
length = np.asanyarray(array).size
1766+
if not (length == N or length == 1):
1767+
raise ValueError(
1768+
f'Argument {name} has a size {length} which does not '
1769+
f'match {N}, the number of ellipses'
1770+
)
1771+
17561772
def _set_transforms(self):
17571773
"""Calculate transforms immediately before drawing."""
17581774

@@ -1796,12 +1812,15 @@ def _set_transforms(self):
17961812
self.set_transform(_affine(m))
17971813

17981814
def _set_widths(self, widths):
1815+
self._check_length(widths, "widths")
17991816
self._widths = 0.5 * np.asarray(widths).ravel()
18001817

18011818
def _set_heights(self, heights):
1819+
self._check_length(heights, "heights")
18021820
self._heights = 0.5 * np.asarray(heights).ravel()
18031821

18041822
def _set_angles(self, angles):
1823+
self._check_length(angles, "angles")
18051824
self._angles = np.deg2rad(angles).ravel()
18061825

18071826
def set_widths(self, widths):

‎lib/matplotlib/tests/test_collections.py

Lines changed: 43 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -428,9 +428,9 @@ def test_EllipseCollection_setter_getter():
428428
offset_transform=ax.transData,
429429
)
430430

431-
assert_array_almost_equal(ec._widths, np.array(widths).ravel() * 0.5)
432-
assert_array_almost_equal(ec._heights, np.array(heights).ravel() * 0.5)
433-
assert_array_almost_equal(ec._angles, np.deg2rad(angles).ravel())
431+
assert_array_almost_equal(ec._widths, np.array(widths) * 0.5)
432+
assert_array_almost_equal(ec._heights, np.array(heights) * 0.5)
433+
assert_array_almost_equal(ec._angles, np.deg2rad(angles))
434434

435435
assert_array_almost_equal(ec.get_widths(), widths)
436436
assert_array_almost_equal(ec.get_heights(), heights)
@@ -440,15 +440,49 @@ def test_EllipseCollection_setter_getter():
440440
ax.set_xlim(-2, 12)
441441
ax.set_ylim(-2, 12)
442442

443-
new_widths = rng.random((10, 2)) * 2
444-
new_heights = rng.random((10, 2)) * 3
445-
new_angles = rng.random((10, 2)) * 180
443+
new_widths = rng.random((10, )) * 2
444+
new_heights = rng.random((10, )) * 3
445+
new_angles = rng.random((10, )) * 180
446446

447447
ec.set(widths=new_widths, heights=new_heights, angles=new_angles)
448448

449-
assert_array_almost_equal(ec.get_widths(), new_widths.ravel())
450-
assert_array_almost_equal(ec.get_heights(), new_heights.ravel())
451-
assert_array_almost_equal(ec.get_angles(), new_angles.ravel())
449+
assert_array_almost_equal(ec.get_widths(), new_widths)
450+
assert_array_almost_equal(ec.get_heights(), new_heights)
451+
assert_array_almost_equal(ec.get_angles(), new_angles)
452+
453+
for array, name in zip(
454+
(new_widths, new_heights, new_angles), ("widths", "heights", "angles")
455+
):
456+
# length doesn't match number of ellipse in the collection
457+
with pytest.raises(ValueError):
458+
kwargs = {name: array[:2]}
459+
ec.set(**kwargs)
460+
461+
# setting a single value
462+
kwargs = {name: array[0]}
463+
ec.set(**kwargs)
464+
assert_array_almost_equal(getattr(ec, f"get_{name}")(), kwargs[name])
465+
466+
467+
def test_EllipseCollection_wrong_length_argument():
468+
# Test widths, heights and angle setter
469+
rng = np.random.default_rng(0)
470+
471+
with pytest.raises(ValueError):
472+
_ = mcollections.EllipseCollection(
473+
widths=(2, 3),
474+
heights=(4, ),
475+
angles=(45, ),
476+
offsets=rng.random((10, 2)) * 10,
477+
)
478+
479+
with pytest.raises(ValueError):
480+
_ = mcollections.EllipseCollection(
481+
widths=(2, ),
482+
heights=(3, 4),
483+
angles=(45, ),
484+
offsets=rng.random((10, 2)) * 10,
485+
)
452486

453487

454488
@image_comparison(['polycollection_close.png'], remove_text=True, style='mpl20')

0 commit comments

Comments
 (0)
0