Skip to content

Commit c89ad62

Browse files
committed
FIX: Accept array for zdir
While we have only documented 3-tuple, it makes sense to accept an array, because a more complicated direction vector could likely be calculated through a rotation matrix. Closes matplotlib#30868.
1 parent b70e324 commit c89ad62

File tree

2 files changed

+20
-3
lines changed

2 files changed

+20
-3
lines changed

lib/mpl_toolkits/mplot3d/art3d.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,11 +58,11 @@ def get_dir_vector(zdir):
5858
x, y, z : array
5959
The direction vector.
6060
"""
61-
if zdir == 'x':
61+
if cbook._str_equal(zdir, 'x'):
6262
return np.array((1, 0, 0))
63-
elif zdir == 'y':
63+
elif cbook._str_equal(zdir, 'y'):
6464
return np.array((0, 1, 0))
65-
elif zdir == 'z':
65+
elif cbook._str_equal(zdir, 'z'):
6666
return np.array((0, 0, 1))
6767
elif zdir is None:
6868
return np.array((0, 0, 0))

lib/mpl_toolkits/mplot3d/tests/test_art3d.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,32 @@
11
import numpy as np
2+
import numpy.testing as nptest
3+
import pytest
24

35
import matplotlib.pyplot as plt
46

57
from matplotlib.backend_bases import MouseEvent
68
from mpl_toolkits.mplot3d.art3d import (
9+
get_dir_vector,
710
Line3DCollection,
811
Poly3DCollection,
912
_all_points_on_plane,
1013
)
1114

1215

16+
@pytest.mark.parametrize("zdir, expected", [
17+
("x", (1, 0, 0)),
18+
("y", (0, 1, 0)),
19+
("z", (0, 0, 1)),
20+
(None, (0, 0, 0)),
21+
((1, 2, 3), (1, 2, 3)),
22+
(np.array([4, 5, 6]), (4, 5, 6)),
23+
])
24+
def test_get_dir_vector(zdir, expected):
25+
res = get_dir_vector(zdir)
26+
assert isinstance(res, np.ndarray)
27+
nptest.assert_array_equal(res, expected)
28+
29+
1330
def test_scatter_3d_projection_conservation():
1431
fig = plt.figure()
1532
ax = fig.add_subplot(projection='3d')

0 commit comments

Comments
 (0)