Add styling support to Check and Radio buttons

QuLogic · QuLogic · commit 2c895fb01600 · 2022-12-30T00:06:56.000-05:00
diff --git a/examples/widgets/check_buttons.py b/examples/widgets/check_buttons.py
@@ -20,9 +20,9 @@
 s2 = np.sin(6*np.pi*t)
 
 fig, ax = plt.subplots()
-l0, = ax.plot(t, s0, visible=False, lw=2, color='k', label='2 Hz')
-l1, = ax.plot(t, s1, lw=2, color='r', label='4 Hz')
-l2, = ax.plot(t, s2, lw=2, color='g', label='6 Hz')
+l0, = ax.plot(t, s0, visible=False, lw=2, color='k', label='1 Hz')
+l1, = ax.plot(t, s1, lw=2, color='r', label='2 Hz')
+l2, = ax.plot(t, s2, lw=2, color='g', label='3 Hz')
 fig.subplots_adjust(left=0.2)
 
 lines_by_label = {l.get_label(): l for l in [l0, l1, l2]}
@@ -32,7 +32,10 @@
 check = CheckButtons(
     ax=rax,
     labels=lines_by_label.keys(),
-    actives=[l.get_visible() for l in lines_by_label.values()]
+    actives=[l.get_visible() for l in lines_by_label.values()],
+    label_props=[{'color': 'k'}, {'color': 'r'}, {'color': 'g'}],
+    frame_props={'edgecolor': ['k', 'r', 'g']},
+    check_props={'facecolor': ['k', 'r', 'g']},
 )
 
 
diff --git a/examples/widgets/radio_buttons.py b/examples/widgets/radio_buttons.py
@@ -25,23 +25,30 @@
 
 axcolor = 'lightgoldenrodyellow'
 rax = fig.add_axes([0.05, 0.7, 0.15, 0.15], facecolor=axcolor)
-radio = RadioButtons(rax, ('2 Hz', '4 Hz', '8 Hz'))
+radio = RadioButtons(rax, ('1 Hz', '2 Hz', '4 Hz'),
+                     radio_props={'s': [16, 32, 64]})
 
 
 def hzfunc(label):
-    hzdict = {'2 Hz': s0, '4 Hz': s1, '8 Hz': s2}
+    hzdict = {'1 Hz': s0, '2 Hz': s1, '4 Hz': s2}
     ydata = hzdict[label]
     l.set_ydata(ydata)
-    plt.draw()
+    fig.canvas.draw()
 radio.on_clicked(hzfunc)
 
 rax = fig.add_axes([0.05, 0.4, 0.15, 0.15], facecolor=axcolor)
-radio2 = RadioButtons(rax, ('red', 'blue', 'green'))
+radio2 = RadioButtons(
+    rax, ('red', 'blue', 'green'),
+    label_props=[{'color': 'red'}, {'color': 'blue'}, {'color': 'green'}],
+    radio_props={
+        'facecolor': ['red', 'blue', 'green'],
+        'edgecolor': ['darkred', 'darkblue', 'darkgreen'],
+    })
 
 
 def colorfunc(label):
     l.set_color(label)
-    plt.draw()
+    fig.canvas.draw()
 radio2.on_clicked(colorfunc)
 
 rax = fig.add_axes([0.05, 0.1, 0.15, 0.15], facecolor=axcolor)
@@ -50,7 +57,7 @@ def colorfunc(label):
 
 def stylefunc(label):
     l.set_linestyle(label)
-    plt.draw()
+    fig.canvas.draw()
 radio3.on_clicked(stylefunc)
 
 plt.show()
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -987,7 +987,8 @@ class CheckButtons(AxesWidget):
         each box, but have ``set_visible(False)`` when its box is not checked.
     """
 
-    def __init__(self, ax, labels, actives=None, *, useblit=True):
+    def __init__(self, ax, labels, actives=None, *, useblit=True,
+                 label_props=None, frame_props=None, check_props=None):
         """
         Add check buttons to `matplotlib.axes.Axes` instance *ax*.
 
@@ -1003,6 +1004,17 @@ def __init__(self, ax, labels, actives=None, *, useblit=True):
         useblit : bool, default: True
             Use blitting for faster drawing if supported by the backend.
             See the tutorial :doc:`/tutorials/advanced/blitting` for details.
+        label_props : dict or list of dict, optional
+            Dictionary of `.Text` properties to be used for the labels. May
+            also be a list of dictionaries, one for each label.
+        frame_props : dict, optional
+            Dictionary of scatter `.Collection` properties to be used for the
+            check button frame. Defaults to 's' marker, (label font size /
+            2)**2 size, black edgecolor, no facecolor, and 1.0 linewidth.
+        check_props : dict, optional
+            Dictionary of scatter `.Collection` properties to be used for the
+            check button check. Defaults to 'x' marker, (label font size /
+            2)**2 size, black color, and 1.0 linewidth.
         """
         super().__init__(ax)
 
@@ -1017,22 +1029,49 @@ def __init__(self, ax, labels, actives=None, *, useblit=True):
         self._background = None
 
         ys = np.linspace(1, 0, len(labels)+2)[1:-1]
-        text_size = mpl.rcParams["font.size"] / 2
 
+        if label_props is None:
+            label_props = [{}] * len(labels)
+        if isinstance(label_props, dict):
+            label_props = [label_props] * len(labels)
         self.labels = [
             ax.text(0.25, y, label, transform=ax.transAxes,
-                    horizontalalignment="left", verticalalignment="center")
-            for y, label in zip(ys, labels)]
-
-        self._squares = ax.scatter(
-            [0.15] * len(ys), ys, marker='s', s=text_size**2,
-            c="none", linewidth=1, transform=ax.transAxes, edgecolor="k"
-        )
-        self._crosses = ax.scatter(
-            [0.15] * len(ys), ys, marker='x', linewidth=1, s=text_size**2,
-            c=["k" if active else "none" for active in actives],
-            transform=ax.transAxes, animated=self._useblit,
-        )
+                    horizontalalignment="left", verticalalignment="center",
+                    **props)
+            for y, label, props in zip(ys, labels, label_props)]
+        text_size = np.array([text.get_fontsize() for text in self.labels]) / 2
+
+        frame_props = {
+            'marker': 's', 's': text_size**2,
+            'facecolor': 'none', 'linewidth': 1, 'edgecolor': 'k',
+            **(frame_props or {}),
+            'transform': ax.transAxes,
+        }
+        self._squares = ax.scatter([0.15] * len(ys), ys, **frame_props)
+        check_props = {
+            'marker': 'x', 'linewidth': 1, 's': text_size**2, 'c': 'k',
+            **(check_props or {}),
+            'transform': ax.transAxes,
+            'animated': self._useblit,
+        }
+        self._checks = ax.scatter([0.15] * len(ys), ys, **check_props)
+        # The user may have passed custom marker and colours in check_props, so
+        # we need to create the checks, and modify the visibility after getting
+        # whatever the user set.
+        self._check_edgecolors = self._checks.get_edgecolor()
+        if len(self._check_edgecolors) == 1:
+            self._check_edgecolors = np.repeat(self._check_edgecolors,
+                                               len(labels), axis=0)
+        self._checks.set_edgecolor(
+            [ec if active else "none"
+             for ec, active in zip(self._check_edgecolors, actives)])
+        self._check_facecolors = self._checks.get_facecolor()
+        if len(self._check_facecolors) == 1:
+            self._check_facecolors = np.repeat(self._check_facecolors,
+                                               len(labels), axis=0)
+        self._checks.set_facecolor(
+            [fc if active else "none"
+             for fc, active in zip(self._check_facecolors, actives)])
 
         self.connect_event('button_press_event', self._clicked)
         if self._useblit:
@@ -1045,7 +1084,7 @@ def _clear(self, event):
         if self.ignore(event):
             return
         self._background = self.canvas.copy_from_bbox(self.ax.bbox)
-        self.ax.draw_artist(self._crosses)
+        self.ax.draw_artist(self._checks)
         if hasattr(self, '_lines'):
             for l1, l2 in self._lines:
                 self.ax.draw_artist(l1)
@@ -1095,15 +1134,25 @@ def set_active(self, index):
         if index not in range(len(self.labels)):
             raise ValueError(f'Invalid CheckButton index: {index}')
 
-        cross_facecolors = self._crosses.get_facecolor()
-        cross_facecolors[index] = colors.to_rgba(
-            "black"
+        check_facecolors = self._checks.get_facecolor()
+        check_facecolors[index] = colors.to_rgba(
+            self._check_facecolors[index]
+            if colors.same_color(
+                check_facecolors[index], colors.to_rgba("none")
+            )
+            else "none"
+        )
+        self._checks.set_facecolor(check_facecolors)
+
+        check_edgecolors = self._checks.get_edgecolor()
+        check_edgecolors[index] = colors.to_rgba(
+            self._check_edgecolors[index]
             if colors.same_color(
-                cross_facecolors[index], colors.to_rgba("none")
+                check_edgecolors[index], colors.to_rgba("none")
             )
             else "none"
         )
-        self._crosses.set_facecolor(cross_facecolors)
+        self._checks.set_edgecolor(check_edgecolors)
 
         if hasattr(self, "_lines"):
             l1, l2 = self._lines[index]
@@ -1114,7 +1163,7 @@ def set_active(self, index):
             if self._useblit:
                 if self._background is not None:
                     self.canvas.restore_region(self._background)
-                self.ax.draw_artist(self._crosses)
+                self.ax.draw_artist(self._checks)
                 if hasattr(self, "_lines"):
                     for l1, l2 in self._lines:
                         self.ax.draw_artist(l1)
@@ -1131,7 +1180,7 @@ def get_status(self):
         Return a list of the status (True/False) of all of the check buttons.
         """
         return [not colors.same_color(color, colors.to_rgba("none"))
-                for color in self._crosses.get_facecolors()]
+                for color in self._checks.get_facecolors()]
 
     def on_clicked(self, func):
         """
@@ -1173,7 +1222,7 @@ def rectangles(self):
     def lines(self):
         if not hasattr(self, "_lines"):
             ys = np.linspace(1, 0, len(self.labels)+2)[1:-1]
-            self._crosses.set_visible(False)
+            self._checks.set_visible(False)
             dy = 1. / (len(self.labels) + 1)
             w, h = dy / 2, dy / 2
             self._lines = []
@@ -1486,7 +1535,7 @@ class RadioButtons(AxesWidget):
     """
 
     def __init__(self, ax, labels, active=0, activecolor='blue', *,
-                 useblit=True):
+                 useblit=True, label_props=None, radio_props=None):
         """
         Add radio buttons to an `~.axes.Axes`.
 
@@ -1503,6 +1552,14 @@ def __init__(self, ax, labels, active=0, activecolor='blue', *,
         useblit : bool, default: True
             Use blitting for faster drawing if supported by the backend.
             See the tutorial :doc:`/tutorials/advanced/blitting` for details.
+        label_props : dict or list of dict, optional
+            Dictionary of `.Text` properties to be used for the labels. May
+            also be a list of dictionaries, one for each label.
+        radio_props : dict, optional
+            Dictionary of scatter `.Collection` properties to be used for the
+            check button frame. Defaults to 'o' marker, (label font size /
+            2)**2 size, black edgecolor, and *activecolor* facecolor (when
+            active).
         """
         super().__init__(ax)
         self.activecolor = activecolor
@@ -1513,19 +1570,41 @@ def __init__(self, ax, labels, active=0, activecolor='blue', *,
         ax.set_navigate(False)
 
         ys = np.linspace(1, 0, len(labels) + 2)[1:-1]
-        text_size = mpl.rcParams["font.size"] / 2
 
         self._useblit = useblit and self.canvas.supports_blit
         self._background = None
 
+        if label_props is None:
+            label_props = [{}] * len(labels)
+        elif isinstance(label_props, dict):
+            label_props = [label_props] * len(labels)
         self.labels = [
             ax.text(0.25, y, label, transform=ax.transAxes,
-                    horizontalalignment="left", verticalalignment="center")
-            for y, label in zip(ys, labels)]
-        self._buttons = ax.scatter(
-            [.15] * len(ys), ys, transform=ax.transAxes, s=text_size**2,
-            c=[activecolor if i == active else "none" for i in range(len(ys))],
-            edgecolor="black", animated=self._useblit)
+                    horizontalalignment="left", verticalalignment="center",
+                    **props)
+            for y, label, props in zip(ys, labels, label_props)]
+        text_size = np.array([text.get_fontsize() for text in self.labels]) / 2
+
+        radio_props = {
+            'marker': 'o',
+            's': text_size**2,
+            'facecolor': activecolor,
+            'edgecolor': 'black',
+            **(radio_props or {}),
+            'transform': ax.transAxes,
+            'animated': self._useblit,
+        }
+        self._buttons = ax.scatter([.15] * len(ys), ys, **radio_props)
+        # The user may have passed custom marker and colours in radio_props, so
+        # we need to create the radios, and modify the visibility after getting
+        # whatever the user set.
+        self._active_colors = self._buttons.get_facecolor()
+        if len(self._active_colors) == 1:
+            self._active_colors = np.repeat(self._active_colors, len(labels),
+                                            axis=0)
+        self._buttons.set_facecolor(
+            [activecolor if i == active else "none"
+             for i, activecolor in enumerate(self._active_colors)])
 
         self.connect_event('button_press_event', self._clicked)
         if self._useblit:
@@ -1576,7 +1655,7 @@ def set_active(self, index):
         self.value_selected = self.labels[index].get_text()
         button_facecolors = self._buttons.get_facecolor()
         button_facecolors[:] = colors.to_rgba("none")
-        button_facecolors[index] = colors.to_rgba(self.activecolor)
+        button_facecolors[index] = colors.to_rgba(self._active_colors[index])
         self._buttons.set_facecolor(button_facecolors)
         if hasattr(self, "_circles"):  # Remove once circles is removed.
             for i, p in enumerate(self._circles):
