Skip to content

Commit

Permalink
ListedColormap: don't pass N colors (#9811)
Browse files Browse the repository at this point in the history
* ListedColormap: don't pass N colors

* fix somewhere else

* fix typing
  • Loading branch information
mathause authored Nov 22, 2024
1 parent 38d5308 commit b18afbf
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 5 deletions.
7 changes: 4 additions & 3 deletions xarray/plot/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,8 +131,9 @@ def _color_palette(cmap, n_colors):

colors_i = np.linspace(0, 1.0, n_colors)
if isinstance(cmap, list | tuple):
# we have a list of colors
cmap = ListedColormap(cmap, N=n_colors)
# expand or truncate the list of colors to n_colors
cmap = list(itertools.islice(itertools.cycle(cmap), n_colors))
cmap = ListedColormap(cmap)
pal = cmap(colors_i)
elif isinstance(cmap, str):
# we have some sort of named palette
Expand All @@ -151,7 +152,7 @@ def _color_palette(cmap, n_colors):
pal = sns.color_palette(cmap, n_colors=n_colors)
except ValueError:
# or maybe we just got a single color as a string
cmap = ListedColormap([cmap], N=n_colors)
cmap = ListedColormap([cmap] * n_colors)
pal = cmap(colors_i)
else:
# cmap better be a LinearSegmentedColormap (e.g. viridis)
Expand Down
5 changes: 3 additions & 2 deletions xarray/tests/test_plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -1805,8 +1805,9 @@ def test_colors_np_levels(self) -> None:
artist = self.darray.plot.contour(levels=levels, colors=["k", "r", "w", "b"])
cmap = artist.cmap
assert isinstance(cmap, mpl.colors.ListedColormap)
colors = cmap.colors
assert isinstance(colors, list)
# non-optimal typing in matplotlib (ArrayLike)
# https://github.com/matplotlib/matplotlib/blob/84464dd085210fb57cc2419f0d4c0235391d97e6/lib/matplotlib/colors.pyi#L133
colors = cast(np.ndarray, cmap.colors)

assert self._color_as_tuple(colors[1]) == (1.0, 0.0, 0.0)
assert self._color_as_tuple(colors[2]) == (1.0, 1.0, 1.0)
Expand Down

0 comments on commit b18afbf

Please sign in to comment.