ENH: add remove method to cbook.Grouper

- also add tests
matplotlib · WeatherGod · Dec 15, 2015 · Dec 15, 2015 · Dec 15, 2015 · Dec 15, 2015
commit 2dfb60b852375df5a8b1421f06881db4c8ead3a4
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -1721,6 +1721,14 @@ def joined(self, a, b):
        except KeyError:
            return False

+    def remove(self, a):
+        self.clean()
+
+        mapping = self._mapping
+        seta = mapping.pop(ref(a), None)
+        if seta is not None:
+            seta.remove(ref(a))
+
    def __iter__(self):
        """
        Iterate over each of the disjoint sets as a list.

diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -1,5 +1,6 @@
 from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
+import itertools

 from matplotlib.externals import six

@@ -376,3 +377,23 @@ def test_step_fails():
                  np.arange(12))
    assert_raises(ValueError, cbook._step_validation,
                  np.arange(12), np.arange(3))
+
+
+def test_grouper():
+    class dummy():
+        pass
+    a, b, c, d, e = objs = [dummy() for j in range(5)]
+    g = cbook.Grouper()
+    g.join(*objs)
+    assert set(list(g)[0]) == set(objs)
+    assert set(g.get_siblings(a)) == set(objs)
+
+    for other in objs[1:]:
+        assert g.joined(a, other)
+
+    g.remove(a)
+    for other in objs[1:]:
+        assert not g.joined(a, other)
+
+    for A, B in itertools.product(objs[1:], objs[1:]):
+        assert g.joined(A, B)