Add fig.mne container for colorbar

mne-tools · Dec 10, 2024 · 7e95b71 · 7e95b71
1 parent b3eb56c
commit 7e95b71
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/mne/viz/circle.py b/mne/viz/circle.py
@@ -6,6 +6,7 @@
 
 from functools import partial
 from itertools import cycle
+from types import SimpleNamespace
 
 import numpy as np
 
@@ -371,6 +372,7 @@ def _plot_connectivity_circle(
         cb_yticks = plt.getp(cb.ax.axes, "yticklabels")
         cb.ax.tick_params(labelsize=fontsize_colorbar)
         plt.setp(cb_yticks, color=textcolor)
+        fig.mne = SimpleNamespace(colorbar=cb)
 
     # Add callback for interaction
     if interactive:

diff --git a/mne/viz/tests/test_circle.py b/mne/viz/tests/test_circle.py
@@ -16,8 +16,10 @@ def test_plot_channel_labels_circle():
     """Test plotting channel labels in a circle."""
     fig, axes = plot_channel_labels_circle(
         dict(brain=["big", "great", "smart"]),
-        colors=dict(big="r", great="y", smart="b"),
+        colors=dict(big="r", great="y", smart="b"), colorbar=True
     )
+    # check that colorbar handle is returned
+    assert isinstance(fig.mne.colorbar, matplotlib.colorbar.Colorbar)
     texts = [
         child.get_text()
         for child in axes.get_children()