Merge pull request #18525 from QuLogic/text3d-simplification

dopplershift · web-flow · commit 5b835eeeba32 · 2020-09-22T20:01:27.000-06:00
Add Text3D position getter/setter
diff --git a/doc/api/next_api_changes/behavior/18525-ES.rst b/doc/api/next_api_changes/behavior/18525-ES.rst
@@ -0,0 +1,5 @@
+``mplot3d.art3d.get_dir_vector`` always returns NumPy arrays
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+For consistency, `~.mplot3d.art3d.get_dir_vector` now always returns NumPy
+arrays, even if the input is a 3-element iterable.
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -12,7 +12,7 @@
 import numpy as np
 
 from matplotlib import (
-    artist, colors as mcolors, lines, text as mtext, path as mpath)
+    artist, cbook, colors as mcolors, lines, text as mtext, path as mpath)
 from matplotlib.collections import (
     LineCollection, PolyCollection, PatchCollection, PathCollection)
 from matplotlib.colors import Normalize
@@ -49,13 +49,12 @@ def get_dir_vector(zdir):
         - 'y': equivalent to (0, 1, 0)
         - 'z': equivalent to (0, 0, 1)
         - *None*: equivalent to (0, 0, 0)
-        - an iterable (x, y, z) is returned unchanged.
+        - an iterable (x, y, z) is converted to a NumPy array, if not already
 
     Returns
     -------
     x, y, z : array-like
-        The direction vector. This is either a numpy.array or *zdir* itself if
-        *zdir* is already a length-3 iterable.
+        The direction vector.
     """
     if zdir == 'x':
         return np.array((1, 0, 0))
@@ -66,7 +65,7 @@ def get_dir_vector(zdir):
     elif zdir is None:
         return np.array((0, 0, 0))
     elif np.iterable(zdir) and len(zdir) == 3:
-        return zdir
+        return np.array(zdir)
     else:
         raise ValueError("'x', 'y', 'z', None or vector of length 3 expected")
 
@@ -95,22 +94,55 @@ def __init__(self, x=0, y=0, z=0, text='', zdir='z', **kwargs):
         mtext.Text.__init__(self, x, y, text, **kwargs)
         self.set_3d_properties(z, zdir)
 
+    def get_position_3d(self):
+        """Return the (x, y, z) position of the text."""
+        return self._x, self._y, self._z
+
+    def set_position_3d(self, xyz, zdir=None):
+        """
+        Set the (*x*, *y*, *z*) position of the text.
+
+        Parameters
+        ----------
+        xyz : (float, float, float)
+            The position in 3D space.
+        zdir : {'x', 'y', 'z', None, 3-tuple}
+            The direction of the text. If unspecified, the zdir will not be
+            changed.
+        """
+        super().set_position(xyz[:2])
+        self.set_z(xyz[2])
+        if zdir is not None:
+            self._dir_vec = get_dir_vector(zdir)
+
+    def set_z(self, z):
+        """
+        Set the *z* position of the text.
+
+        Parameters
+        ----------
+        z : float
+        """
+        self._z = z
+        self.stale = True
+
     def set_3d_properties(self, z=0, zdir='z'):
-        x, y = self.get_position()
-        self._position3d = np.array((x, y, z))
+        self._z = z
         self._dir_vec = get_dir_vector(zdir)
         self.stale = True
 
     @artist.allow_rasterization
     def draw(self, renderer):
+        position3d = np.array((self._x, self._y, self._z))
         proj = proj3d.proj_trans_points(
-            [self._position3d, self._position3d + self._dir_vec], renderer.M)
+            [position3d, position3d + self._dir_vec],
+            renderer.M)
         dx = proj[0][1] - proj[0][0]
         dy = proj[1][1] - proj[1][0]
         angle = math.degrees(math.atan2(dy, dx))
-        self.set_position((proj[0][0], proj[1][0]))
-        self.set_rotation(_norm_text_angle(angle))
-        mtext.Text.draw(self, renderer)
+        with cbook._setattr_cm(self, _x=proj[0][0], _y=proj[1][0],
+                               _rotation=_norm_text_angle(angle)):
+            mtext.Text.draw(self, renderer)
         self.stale = False
 
     def get_tightbbox(self, renderer):
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -416,6 +416,31 @@ def test_text3d():
     ax.set_zlabel('Z axis')
 
 
+@check_figures_equal(extensions=['png'])
+def test_text3d_modification(fig_ref, fig_test):
+    # Modifying the Text position after the fact should work the same as
+    # setting it directly.
+    zdirs = (None, 'x', 'y', 'z', (1, 1, 0), (1, 1, 1))
+    xs = (2, 6, 4, 9, 7, 2)
+    ys = (6, 4, 8, 7, 2, 2)
+    zs = (4, 2, 5, 6, 1, 7)
+
+    ax_test = fig_test.add_subplot(projection='3d')
+    ax_test.set_xlim3d(0, 10)
+    ax_test.set_ylim3d(0, 10)
+    ax_test.set_zlim3d(0, 10)
+    for zdir, x, y, z in zip(zdirs, xs, ys, zs):
+        t = ax_test.text(0, 0, 0, f'({x}, {y}, {z}), dir={zdir}')
+        t.set_position_3d((x, y, z), zdir=zdir)
+
+    ax_ref = fig_ref.add_subplot(projection='3d')
+    ax_ref.set_xlim3d(0, 10)
+    ax_ref.set_ylim3d(0, 10)
+    ax_ref.set_zlim3d(0, 10)
+    for zdir, x, y, z in zip(zdirs, xs, ys, zs):
+        ax_ref.text(x, y, z, f'({x}, {y}, {z}), dir={zdir}', zdir=zdir)
+
+
 @mpl3d_image_comparison(['trisurf3d.png'], tol=0.03)
 def test_trisurf3d():
     n_angles = 36
