Make Grouper return siblings in the order in which they have been seen.

anntzer · anntzer · commit 63a5d254e687 · 2023-03-27T10:45:26.000+02:00
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -3911,19 +3911,13 @@ def format_coord(self, x, y):
             return "x={} y={}".format(
                 "???" if x is None else self.format_xdata(x),
                 "???" if y is None else self.format_ydata(y))
-        x_twins = []
-        y_twins = []
-        # Retrieve twins in the order of self.figure.axes to sort tied zorders (which is
-        # the common case) by the order in which they are added to the figure.
-        for ax in self.figure.axes:
-            if ax in twins:
-                if ax is self or ax._sharex is self or self._sharex is ax:
-                    x_twins.append(ax)
-                if ax is self or ax._sharey is self or self._sharey is ax:
-                    y_twins.append(ax)
         get_zorder = attrgetter("zorder")
-        x_twins.sort(key=get_zorder)
-        y_twins.sort(key=get_zorder)
+        x_twins = sorted([ax for ax in twins
+                          if ax is self or ax._sharex is self or self._sharex is ax],
+                         key=get_zorder)
+        y_twins = sorted([ax for ax in twins
+                          if ax is self or ax._sharey is self or self._sharey is ax],
+                         key=get_zorder)
         screen_xy = self.transData.transform((x, y))
         return "x={} y={}".format(
             " / ".join(
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -783,21 +783,27 @@ class Grouper:
     """
 
     def __init__(self, init=()):
-        self._mapping = weakref.WeakKeyDictionary(
-            {x: weakref.WeakSet([x]) for x in init})
+        self._count = itertools.count()
+        # For each item, we store (order_in_which_item_was_seen, group_of_item), which
+        # lets __iter__ and get_siblings return items in the order in which they have
+        # been seen.
+        self._mapping = weakref.WeakKeyDictionary()
+        for x in init:
+            if x not in self._mapping:
+                self._mapping[x] = (next(self._count), weakref.WeakSet([x]))
 
     def __getstate__(self):
         return {
             **vars(self),
             # Convert weak refs to strong ones.
-            "_mapping": {k: set(v) for k, v in self._mapping.items()},
+            "_mapping": {k: (i, set(v)) for k, (i, v) in self._mapping.items()},
         }
 
     def __setstate__(self, state):
         vars(self).update(state)
         # Convert strong refs to weak ones.
         self._mapping = weakref.WeakKeyDictionary(
-            {k: weakref.WeakSet(v) for k, v in self._mapping.items()})
+            {k: (i, weakref.WeakSet(v)) for k, (i, v) in self._mapping.items()})
 
     def __contains__(self, item):
         return item in self._mapping
@@ -810,24 +816,30 @@ def join(self, a, *args):
         """
         Join given arguments into the same set.  Accepts one or more arguments.
         """
-        mapping = self._mapping
-        set_a = mapping.setdefault(a, weakref.WeakSet([a]))
-
-        for arg in args:
-            set_b = mapping.get(arg, weakref.WeakSet([arg]))
+        m = self._mapping
+        try:
+            _, set_a = m[a]
+        except KeyError:
+            _, set_a = m.setdefault(a, (next(self._count), weakref.WeakSet([a])))
+        for b in args:
+            try:
+                _, set_b = m[b]
+            except KeyError:
+                _, set_b = m.setdefault(b, (next(self._count), weakref.WeakSet([b])))
             if set_b is not set_a:
                 if len(set_b) > len(set_a):
                     set_a, set_b = set_b, set_a
                 set_a.update(set_b)
                 for elem in set_b:
-                    mapping[elem] = set_a
+                    i, _ = m[elem]
+                    m[elem] = (i, set_a)
 
     def joined(self, a, b):
         """Return whether *a* and *b* are members of the same set."""
-        return (self._mapping.get(a, object()) is self._mapping.get(b))
+        return (self._mapping.get(a, (None, object()))[1] is self._mapping.get(b)[1])
 
     def remove(self, a):
-        set_a = self._mapping.pop(a, None)
+        _, set_a = self._mapping.pop(a, (None, None))
         if set_a:
             set_a.remove(a)
 
@@ -837,14 +849,14 @@ def __iter__(self):
 
         The iterator is invalid if interleaved with calls to join().
         """
-        unique_groups = {id(group): group for group in self._mapping.values()}
+        unique_groups = {id(group): group for _, group in self._mapping.values()}
         for group in unique_groups.values():
-            yield [x for x in group]
+            yield sorted(group, key=self._mapping.__getitem__)
 
     def get_siblings(self, a):
         """Return all of the items joined with *a*, including itself."""
-        siblings = self._mapping.get(a, [a])
-        return [x for x in siblings]
+        _, siblings = self._mapping.get(a, (None, [a]))
+        return sorted(siblings, key=self._mapping.get)
 
 
 class GrouperView:
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -591,9 +591,9 @@ class Dummy:
     for o in objs:
         assert o in mapping
 
-    base_set = mapping[objs[0]]
+    _, base_set = mapping[objs[0]]
     for o in objs[1:]:
-        assert mapping[o] is base_set
+        assert mapping[o][1] is base_set
 
 
 def test_flatiter():
