ENH: add remove method to cbook.Grouper

tacaswell · tacaswell · commit 2dfb60b85237 · 2015-12-14T23:52:06.000-05:00
- also add tests
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)
