Add getter/setter for Text3D position.

QuLogic · QuLogic · commit 7034b70b3b90 · 2020-09-22T20:54:00.000-04:00
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -94,6 +94,38 @@ 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'):
         self._z = z
         self._dir_vec = get_dir_vector(zdir)
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
