Path3DCollection restore vert orders after draw() and index conservation test

hukaidong · hukaidong · commit 94e21bd577aa · 2022-11-18T18:18:09.000-05:00
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -11,6 +11,8 @@
 
 import numpy as np
 
+from contextlib import contextmanager
+
 from matplotlib import (
     artist, cbook, colors as mcolors, lines, text as mtext,
     path as mpath)
@@ -629,10 +631,12 @@ def __init__(self, *args, zs=0, zdir='z', depthshade=True, **kwargs):
         self._in_draw = False
         super().__init__(*args, **kwargs)
         self.set_3d_properties(zs, zdir)
+        self._offset_zordered = None
 
     def draw(self, renderer):
-        with cbook._setattr_cm(self, _in_draw=True):
-            super().draw(renderer)
+        with self._use_zordered_offset():
+            with cbook._setattr_cm(self, _in_draw=True):
+                super().draw(renderer)
 
     def set_sort_zpos(self, val):
         """Set the position to use for z-sorting."""
@@ -731,15 +735,32 @@ def do_3d_projection(self):
         if len(self._linewidths3d) > 1:
             self._linewidths = self._linewidths3d[z_markers_idx]
 
+        PathCollection.set_offsets(self, np.column_stack((vxs, vys)))
+
         # Re-order items
         vzs = vzs[z_markers_idx]
         vxs = vxs[z_markers_idx]
         vys = vys[z_markers_idx]
 
-        PathCollection.set_offsets(self, np.column_stack((vxs, vys)))
+        # Store ordered offset for drawing purpose
+        self._offset_zordered = np.column_stack((vxs, vys))
 
         return np.min(vzs) if vzs.size else np.nan
 
+    @contextmanager
+    def _use_zordered_offset(self):
+        if self._offset_zordered is None:
+            # Do nothing
+            yield
+        else:
+            # Swap offset with z-ordered offset
+            old_offset = self._offsets
+            super().set_offsets(self._offset_zordered)
+            try:
+                yield
+            finally:
+                self._offsets = old_offset
+
     def _maybe_depth_shade_and_sort_colors(self, color_array):
         color_array = (
             _zalpha(color_array, self._vzs)
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_art3d.py b/lib/mpl_toolkits/mplot3d/tests/test_art3d.py
@@ -0,0 +1,38 @@
+import matplotlib.pyplot as plt
+
+from matplotlib.backend_bases import MouseEvent
+
+
+def test_scatter_3d_projection_conservation():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+    # fix axes3d projection
+    ax.roll = 0
+    ax.elev = 0
+    ax.azim = -45
+    ax.stale = True
+
+    x = [0, 1, 2, 3, 4]
+    scatter_collection = ax.scatter(x, x, x)
+    fig.canvas.draw_idle()
+
+    # Get scatter location on canvas and freeze the data
+    scatter_offset = scatter_collection.get_offsets()
+    scatter_location = ax.transData.transform(scatter_offset)
+
+    # Yaw -44 and -46 are enough to produce two set of scatter
+    # with opposite z-order without moving points too far
+    for azim in (-44, -46):
+        ax.azim = azim
+        ax.stale = True
+        fig.canvas.draw_idle()
+
+        for i in range(5):
+            # Create a mouse event used to locate and to get index
+            # from each dots
+            event = MouseEvent("button_press_event", fig.canvas,
+                               *scatter_location[i, :])
+            contains, ind = scatter_collection.contains(event)
+            assert contains is True
+            assert len(ind["ind"]) == 1
+            assert ind["ind"][0] == i
