matplotlib
diff --git a/‎lib/mpl_toolkits/mplot3d/art3d.py
Lines changed: 6 additions & 1 deletion b/‎lib/mpl_toolkits/mplot3d/art3d.py
Lines changed: 6 additions & 1 deletion
diff --git a/‎lib/mpl_toolkits/mplot3d/axes3d.py
Lines changed: 2 additions & 1 deletion b/‎lib/mpl_toolkits/mplot3d/axes3d.py
Lines changed: 2 additions & 1 deletion
diff --git a/‎lib/mpl_toolkits/tests/baseline_images/test_mplot3d/add_collection3d_zs_array.png
59.2 KB b/‎lib/mpl_toolkits/tests/baseline_images/test_mplot3d/add_collection3d_zs_array.png
59.2 KB
diff --git a/‎lib/mpl_toolkits/tests/baseline_images/test_mplot3d/add_collection3d_zs_scalar.png
50.5 KB b/‎lib/mpl_toolkits/tests/baseline_images/test_mplot3d/add_collection3d_zs_scalar.png
50.5 KB
diff --git a/‎lib/mpl_toolkits/tests/test_mplot3d.py
Lines changed: 54 additions & 0 deletions b/‎lib/mpl_toolkits/tests/test_mplot3d.py
Lines changed: 54 additions & 0 deletions
@@ -207,7 +207,12 @@ def _path_to_3d_segment(path, zs=0, zdir='z'):
 def _paths_to_3d_segments(paths, zs=0, zdir='z'):
     """Convert paths from a collection object to 3D segments."""
 
-    zs = np.broadcast_to(zs, len(paths))
+    if not np.iterable(zs):
+        zs = np.broadcast_to(zs, len(paths))
+    else:
+        if len(zs) != len(paths):
+            raise ValueError('Number of z-coordinates does not match paths.')
+
     segs = [_path_to_3d_segment(path, pathz, zdir)
             for path, pathz in zip(paths, zs)]
     return segs
 
@@ -2250,7 +2250,8 @@ def add_collection3d(self, col, zs=0, zdir='z'):
             art3d.patch_collection_2d_to_3d(col, zs=zs, zdir=zdir)
             col.set_sort_zpos(zsortval)
 
-        super().add_collection(col)
+        collection = super().add_collection(col)
+        return collection
 
     def scatter(self, xs, ys, zs=0, zdir='z', s=20, c=None, depthshade=True,
                 *args, **kwargs):
 
@@ -543,6 +543,60 @@ def test_poly3dcollection_alpha():
     ax.add_collection3d(c2)
 
 
+@mpl3d_image_comparison(['add_collection3d_zs_array.png'])
+def test_add_collection3d_zs_array():
+    theta = np.linspace(-4 * np.pi, 4 * np.pi, 100)
+    z = np.linspace(-2, 2, 100)
+    r = z**2 + 1
+    x = r * np.sin(theta)
+    y = r * np.cos(theta)
+
+    points = np.column_stack([x, y, z]).reshape(-1, 1, 3)
+    segments = np.concatenate([points[:-1], points[1:]], axis=1)
+
+    fig = plt.figure()
+    ax = fig.gca(projection='3d')
+
+    norm = plt.Normalize(0, 2*np.pi)
+    # 2D LineCollection from x & y values
+    lc = LineCollection(segments[:, :, :2], cmap='twilight', norm=norm)
+    lc.set_array(np.mod(theta, 2*np.pi))
+    # Add 2D collection at z values to ax
+    line = ax.add_collection3d(lc, zs=segments[:, :, 2])
+
+    assert line is not None
+
+    ax.set_xlim(-5, 5)
+    ax.set_ylim(-4, 6)
+    ax.set_zlim(-2, 2)
+
+
+@mpl3d_image_comparison(['add_collection3d_zs_scalar.png'])
+def test_add_collection3d_zs_scalar():
+    theta = np.linspace(0, 2 * np.pi, 100)
+    z = 1
+    r = z**2 + 1
+    x = r * np.sin(theta)
+    y = r * np.cos(theta)
+
+    points = np.column_stack([x, y]).reshape(-1, 1, 2)
+    segments = np.concatenate([points[:-1], points[1:]], axis=1)
+
+    fig = plt.figure()
+    ax = fig.gca(projection='3d')
+
+    norm = plt.Normalize(0, 2*np.pi)
+    lc = LineCollection(segments, cmap='twilight', norm=norm)
+    lc.set_array(theta)
+    line = ax.add_collection3d(lc, zs=z)
+
+    assert line is not None
+
+    ax.set_xlim(-5, 5)
+    ax.set_ylim(-4, 6)
+    ax.set_zlim(0, 2)
+
+
 @mpl3d_image_comparison(['axes3d_labelpad.png'], remove_text=False)
 def test_axes3d_labelpad():
     fig = plt.figure()