ENG: add a by_label accessor to FigureRegistry

Author: tacaswell

docs/source/api.rst

@@ -49,6 +49,7 @@ Managed
    mpl_gui.FigureRegistry.figure
    mpl_gui.FigureRegistry.subplots
    mpl_gui.FigureRegistry.subplot_mosaic
+   mpl_gui.FigureRegistry.by_label
    mpl_gui.FigureRegistry.show_all
    mpl_gui.FigureRegistry.close_all
 

docs/source/index.rst

@@ -156,6 +156,19 @@ to ::
 
 and have a (mostly) drop-in replacement.
 
+Additionally, there is a  `.FigureRegistry.by_label` accessory that returns
+a dictionary mapping the Figures' labels to each Figure ::
+
+  import mpl_gui as mg
+
+  fr = mg.FigureRegistry()
+
+  figA = fr.figure(label='A')
+  figB = fr.subplots(2, 2, label='B')
+
+  fr.by_label['A'] is figA
+  fr.by_label['B'] is figB
+
 FigureContext
 +++++++++++++
 

mpl_gui/__init__.py

@@ -14,6 +14,7 @@
 """
 import logging
 import functools
+from itertools import count
 
 from matplotlib.backend_bases import FigureCanvasBase as _FigureCanvasBase
 
@@ -113,19 +114,47 @@ class FigureRegistry:
 
     """
 
-    def __init__(self, *, block=None, timeout=0):
+    def __init__(self, *, block=None, timeout=0, prefix="Figure "):
+        # settings stashed to set defaults on show
         self._timeout = timeout
         self._block = block
+        # Settings / state to control the default figure label
+        self._count = count()
+        self._prefix = prefix
+        # the canonical location for storing the Figures this registry owns.
+        # any additional views must never include a figure not in the list but
+        # may omit figures
         self.figures = []
 
     def _register_fig(self, fig):
+        # if the user closes the figure by any other mechanism, drop our
+        # reference to it.  This is important for getting a "pyplot" like user
+        # experience
         fig.canvas.mpl_connect(
             "close_event",
             lambda e: self.figures.remove(fig) if fig in self.figures else None,
         )
+        # hold a hard reference to the figure.
         self.figures.append(fig)
+        # Make sure we give the figure a quasi-unique label.  We will never set
+        # the same label twice, but will not over-ride any user label (but
+        # empty string) on a Figure so if they provide duplicate labels, change
+        # the labels under us, or provide a label that will be shadowed in the
+        # future it will be what it is.
+        fignum = next(self._count)
+        if fig.get_label() == "":
+            fig.set_label(f"{self._prefix}{fignum:d}")
         return fig
 
+    @property
+    def by_label(self):
+        """
+        Return a dictionary of the current mapping labels -> figures.
+
+        If there are duplicate labels, newer figures will take precedence.
+        """
+        return {fig.get_label(): fig for fig in self.figures}
+
     @functools.wraps(figure)
     def figure(self, *args, **kwargs):
         fig = figure(*args, **kwargs)

mpl_gui/tests/test_examples.py

@@ -7,7 +7,6 @@
 import mpl_gui as mg
 
 
-
 def test_no_pyplot():
 
     assert sys.modules.get("matplotlib.pyplot", None) is None
@@ -79,6 +78,7 @@ class TestException(Exception):
 
         assert isinstance(fig.canvas, FigureCanvasBase)
 
+
 def test_close_all():
     fr = mg.FigureRegistry(block=False)
     fig = fr.figure()
@@ -90,11 +90,51 @@ def test_close_all():
     new_canvas = fig.canvas
     fr.close_all()
     assert len(fr.figures) == 0
-    assert 'destroy' in new_canvas.manager.call_info
+    assert "destroy" in new_canvas.manager.call_info
     assert fig.canvas is not new_canvas
     assert new_canvas.figure is None
 
     # test revive
     old_canvas = fig.canvas
     mg.show([fig])
     assert fig.canvas is not old_canvas
+
+
+def test_labels_prefix():
+    fr = mg.FigureRegistry(block=False, prefix="Aardvark ")
+    for j in range(5):
+        fr.figure()
+        assert list(fr.by_label) == [f"Aardvark {k}" for k in range(j + 1)]
+    fr.close_all()
+    assert len(fr.by_label) == 0
+
+
+def test_labels_collision():
+    fr = mg.FigureRegistry(block=False)
+    for j in range(5):
+        fr.figure(label="aardvark")
+    assert list(fr.by_label) == ["aardvark"]
+    assert len(fr.figures) == 5
+    assert len(set(fr.figures)) == 5
+    assert fr.figures[-1] is fr.by_label["aardvark"]
+    fr.close_all()
+    assert len(fr.by_label) == 0
+
+
+def test_by_label_new_dict():
+    fr = mg.FigureRegistry(block=False)
+    for j in range(5):
+        fr.figure()
+    # test we get a new dict each time!
+    assert fr.by_label is not fr.by_label
+
+
+def test_change_labels():
+    fr = mg.FigureRegistry(block=False)
+    for j in range(5):
+        fr.figure()
+    assert list(fr.by_label) == [f"Figure {j}" for j in range(5)]
+
+    for j, f in enumerate(fr.by_label.values()):
+        f.set_label(f"aardvark {j}")
+    assert list(fr.by_label) == [f"aardvark {j}" for j in range(5)]