Backport PR #30869: FIX: Accept array for zdir (#30870)

meeseeksmachine · rcomer · web-flow · commit 97c4e0734bcb · 2025-12-17T12:06:17.000+01:00
Co-authored-by: Ruth Comer &lt;10599679+rcomer@users.noreply.github.com&gt;
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -59,11 +59,11 @@ def get_dir_vector(zdir):
     x, y, z : array
         The direction vector.
     """
-    if zdir == 'x':
+    if cbook._str_equal(zdir, 'x'):
         return np.array((1, 0, 0))
-    elif zdir == 'y':
+    elif cbook._str_equal(zdir, 'y'):
         return np.array((0, 1, 0))
-    elif zdir == 'z':
+    elif cbook._str_equal(zdir, 'z'):
         return np.array((0, 0, 1))
     elif zdir is None:
         return np.array((0, 0, 0))
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_art3d.py b/lib/mpl_toolkits/mplot3d/tests/test_art3d.py
@@ -1,15 +1,32 @@
 import numpy as np
+import numpy.testing as nptest
+import pytest
 
 import matplotlib.pyplot as plt
 
 from matplotlib.backend_bases import MouseEvent
 from mpl_toolkits.mplot3d.art3d import (
+    get_dir_vector,
     Line3DCollection,
     Poly3DCollection,
     _all_points_on_plane,
 )
 
 
+@pytest.mark.parametrize("zdir, expected", [
+    ("x", (1, 0, 0)),
+    ("y", (0, 1, 0)),
+    ("z", (0, 0, 1)),
+    (None, (0, 0, 0)),
+    ((1, 2, 3), (1, 2, 3)),
+    (np.array([4, 5, 6]), (4, 5, 6)),
+])
+def test_get_dir_vector(zdir, expected):
+    res = get_dir_vector(zdir)
+    assert isinstance(res, np.ndarray)
+    nptest.assert_array_equal(res, expected)
+
+
 def test_scatter_3d_projection_conservation():
     fig = plt.figure()
     ax = fig.add_subplot(projection='3d')
