Merge pull request #10111 from Raab70/bug-fix/line3d

tacaswell · web-flow · commit c044780fdc48 · 2019-01-08T09:44:42.000-05:00
ENH: Add set_data_3d and get_data_3d to Line3d
diff --git a/doc/users/next_whats_new/2018_01_02_line3d_get_set_data.rst b/doc/users/next_whats_new/2018_01_02_line3d_get_set_data.rst
@@ -0,0 +1,7 @@
+mplot3d Line3D now allows {set,get}_data_3d
+-------------------------------------------
+
+Lines created with the 3d projection in mplot3d can now access the data using
+``mplot3d.art3d.Line3D.get_data_3d()`` which returns a tuple of array_likes containing
+the (x, y, z) data. The equivalent ``mplot3d.art3d.Line3D.set_data_3d(x, y, z)``
+can be used to modify the data of an existing Line3D.
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -150,6 +150,40 @@ def set_3d_properties(self, zs=0, zdir='z'):
         self._verts3d = juggle_axes(xs, ys, zs, zdir)
         self.stale = True
 
+    def set_data_3d(self, *args):
+        """
+        Set the x, y and z data
+
+        Parameters
+        ----------
+        x : array_like
+            The x-data to be plotted
+        y : array_like
+            The y-data to be plotted
+        z : array_like
+            The z-data to be plotted
+
+        Notes
+        -----
+        Accepts x, y, z arguments or a single array_like (x, y, z)
+        """
+        if len(args) == 1:
+            self._verts3d = args[0]
+        else:
+            self._verts3d = args
+        self.stale = True
+
+    def get_data_3d(self):
+        """
+        Get the current data
+
+        Returns
+        -------
+        verts3d : length-3 tuple or array_likes
+            The current data as a tuple or array_likes
+        """
+        return self._verts3d
+
     @artist.allow_rasterization
     def draw(self, renderer):
         xs3d, ys3d, zs3d = self._verts3d
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -815,6 +815,18 @@ def test_calling_conventions(self):
             ax.voxels(filled=filled, x=x, y=y, z=z)
 
 
+def test_line3d_set_get_data_3d():
+    x, y, z = [0, 1], [2, 3], [4, 5]
+    x2, y2, z2 = [6, 7], [8, 9], [10, 11]
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+    lines = ax.plot(x, y, z)
+    line = lines[0]
+    np.testing.assert_array_equal((x, y, z), line.get_data_3d())
+    line.set_data_3d(x2, y2, z2)
+    np.testing.assert_array_equal((x2, y2, z2), line.get_data_3d())
+
+
 def test_inverted_cla():
     # Github PR #5450. Setting autoscale should reset
     # axes to be non-inverted.
