DRY up

Author: cbrnr

mne/viz/epochs.py

@@ -29,6 +29,7 @@
     _handle_precompute,
     _make_combine_callable,
     _make_event_color_dict,
+    _normalize_annotation_colors,
     _set_title_multiple_electrodes,
     _set_window_title,
     _setup_cmap,
@@ -869,7 +870,7 @@ def plot_epochs(
     annotation_colors : dict | None
         A dictionary mapping annotation description strings to colors. Use this to
         override the default color assigned to specific annotation types (e.g.,
-        ``dict(bad_segment='orange')``). Colors can be any valid matplotlib color
+        ``dict(bad_segment='orange')``). Colors can be any valid Matplotlib color
         specification. Keys that do not match any annotation description in the data
         will trigger a warning. If ``None`` (default), automatic colors are used.
 
@@ -1026,24 +1027,9 @@ def plot_epochs(
 
     # handle annotation_colors
     if annotation_colors is not None:
-        from matplotlib.colors import to_hex
-
-        _validate_type(annotation_colors, dict, "annotation_colors")
-        normalized = {}
-        for k, v in annotation_colors.items():
-            try:
-                normalized[k] = to_hex(v)
-            except ValueError:
-                raise ValueError(
-                    f"annotation_colors[{k!r}] is not a valid matplotlib color: {v!r}"
-                ) from None
-        unknown = set(normalized) - set(epochs.annotations.description)
-        if unknown:
-            warn(
-                "The following annotation_colors keys do not match any annotation "
-                f"description in the data: {sorted(unknown)}"
-            )
-        annotation_colors = normalized
+        annotation_colors = _normalize_annotation_colors(
+            annotation_colors, epochs.annotations
+        )
 
     params = dict(
         inst=epochs,

mne/viz/raw.py

@@ -27,6 +27,7 @@
     _handle_decim,
     _handle_precompute,
     _make_event_color_dict,
+    _normalize_annotation_colors,
     _shorten_path_from_middle,
 )
 
@@ -115,7 +116,7 @@ def plot_raw(
     annotation_colors : dict | None
         A dictionary mapping annotation description strings to colors. Use this to
         override the default color assigned to specific annotation types (e.g.,
-        ``dict(bad_segment='orange')``). Colors can be any valid matplotlib color
+        ``dict(bad_segment='orange')``). Colors can be any valid Matplotlib color
         specification. Keys that do not match any annotation description in the data
         will trigger a warning. If ``None`` (default), automatic colors are used.
 
@@ -353,24 +354,9 @@ def plot_raw(
 
     # handle annotation_colors
     if annotation_colors is not None:
-        from matplotlib.colors import to_hex
-
-        _validate_type(annotation_colors, dict, "annotation_colors")
-        normalized = {}
-        for k, v in annotation_colors.items():
-            try:
-                normalized[k] = to_hex(v)
-            except ValueError:
-                raise ValueError(
-                    f"annotation_colors[{k!r}] is not a valid matplotlib color: {v!r}"
-                ) from None
-        unknown = set(normalized) - set(raw.annotations.description)
-        if unknown:
-            warn(
-                "The following annotation_colors keys do not match any "
-                f"annotation description in the data: {sorted(unknown)}"
-            )
-        annotation_colors = normalized
+        annotation_colors = _normalize_annotation_colors(
+            annotation_colors, raw.annotations
+        )
 
     # handle event colors
     event_color_dict = _make_event_color_dict(event_color, events, event_id)

mne/viz/utils.py

@@ -2866,3 +2866,34 @@ def _get_plot_ch_type(inst, ch_type, allow_ref_meg=False):
                 f"No plottable channel types found. Allowed types are: {allowed_types}"
             )
     return ch_type
+
+
+def _normalize_annotation_colors(annotation_colors, annotations):
+    """Normalize annotation_colors and check that keys match annotation descriptions.
+
+    Parameters
+    ----------
+    annotation_colors : dict[str, color]
+        The annotation colors to normalize (``color`` can be any valid Matplotlib color
+        specification).
+    annotations : mne.Annotations
+        The Annotations object to check against.
+    """
+    from matplotlib.colors import to_hex
+
+    _validate_type(annotation_colors, dict, "annotation_colors")
+    normalized = {}
+    for k, v in annotation_colors.items():
+        try:
+            normalized[k] = to_hex(v)
+        except ValueError:
+            raise ValueError(
+                f"annotation_colors[{k!r}] is not a valid matplotlib color: {v!r}"
+            ) from None
+    unknown = set(normalized) - set(annotations.description)
+    if unknown:
+        warn(
+            "The following annotation_colors keys do not match any annotation "
+            f"description in the data: {sorted(unknown)}"
+        )
+    return normalized