279 lines
8.2 KiB
Python
279 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import (
|
|
Collection,
|
|
Iterator,
|
|
)
|
|
import itertools
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
cast,
|
|
)
|
|
import warnings
|
|
|
|
import matplotlib as mpl
|
|
import matplotlib.colors
|
|
import numpy as np
|
|
|
|
from pandas._typing import MatplotlibColor as Color
|
|
from pandas.util._exceptions import find_stack_level
|
|
|
|
from pandas.core.dtypes.common import is_list_like
|
|
|
|
import pandas.core.common as com
|
|
|
|
if TYPE_CHECKING:
|
|
from matplotlib.colors import Colormap
|
|
|
|
|
|
def get_standard_colors(
|
|
num_colors: int,
|
|
colormap: Colormap | None = None,
|
|
color_type: str = "default",
|
|
color: dict[str, Color] | Color | Collection[Color] | None = None,
|
|
):
|
|
"""
|
|
Get standard colors based on `colormap`, `color_type` or `color` inputs.
|
|
|
|
Parameters
|
|
----------
|
|
num_colors : int
|
|
Minimum number of colors to be returned.
|
|
Ignored if `color` is a dictionary.
|
|
colormap : :py:class:`matplotlib.colors.Colormap`, optional
|
|
Matplotlib colormap.
|
|
When provided, the resulting colors will be derived from the colormap.
|
|
color_type : {"default", "random"}, optional
|
|
Type of colors to derive. Used if provided `color` and `colormap` are None.
|
|
Ignored if either `color` or `colormap` are not None.
|
|
color : dict or str or sequence, optional
|
|
Color(s) to be used for deriving sequence of colors.
|
|
Can be either be a dictionary, or a single color (single color string,
|
|
or sequence of floats representing a single color),
|
|
or a sequence of colors.
|
|
|
|
Returns
|
|
-------
|
|
dict or list
|
|
Standard colors. Can either be a mapping if `color` was a dictionary,
|
|
or a list of colors with a length of `num_colors` or more.
|
|
|
|
Warns
|
|
-----
|
|
UserWarning
|
|
If both `colormap` and `color` are provided.
|
|
Parameter `color` will override.
|
|
"""
|
|
if isinstance(color, dict):
|
|
return color
|
|
|
|
colors = _derive_colors(
|
|
color=color,
|
|
colormap=colormap,
|
|
color_type=color_type,
|
|
num_colors=num_colors,
|
|
)
|
|
|
|
return list(_cycle_colors(colors, num_colors=num_colors))
|
|
|
|
|
|
def _derive_colors(
|
|
*,
|
|
color: Color | Collection[Color] | None,
|
|
colormap: str | Colormap | None,
|
|
color_type: str,
|
|
num_colors: int,
|
|
) -> list[Color]:
|
|
"""
|
|
Derive colors from either `colormap`, `color_type` or `color` inputs.
|
|
|
|
Get a list of colors either from `colormap`, or from `color`,
|
|
or from `color_type` (if both `colormap` and `color` are None).
|
|
|
|
Parameters
|
|
----------
|
|
color : str or sequence, optional
|
|
Color(s) to be used for deriving sequence of colors.
|
|
Can be either be a single color (single color string, or sequence of floats
|
|
representing a single color), or a sequence of colors.
|
|
colormap : :py:class:`matplotlib.colors.Colormap`, optional
|
|
Matplotlib colormap.
|
|
When provided, the resulting colors will be derived from the colormap.
|
|
color_type : {"default", "random"}, optional
|
|
Type of colors to derive. Used if provided `color` and `colormap` are None.
|
|
Ignored if either `color` or `colormap`` are not None.
|
|
num_colors : int
|
|
Number of colors to be extracted.
|
|
|
|
Returns
|
|
-------
|
|
list
|
|
List of colors extracted.
|
|
|
|
Warns
|
|
-----
|
|
UserWarning
|
|
If both `colormap` and `color` are provided.
|
|
Parameter `color` will override.
|
|
"""
|
|
if color is None and colormap is not None:
|
|
return _get_colors_from_colormap(colormap, num_colors=num_colors)
|
|
elif color is not None:
|
|
if colormap is not None:
|
|
warnings.warn(
|
|
"'color' and 'colormap' cannot be used simultaneously. Using 'color'",
|
|
stacklevel=find_stack_level(),
|
|
)
|
|
return _get_colors_from_color(color)
|
|
else:
|
|
return _get_colors_from_color_type(color_type, num_colors=num_colors)
|
|
|
|
|
|
def _cycle_colors(colors: list[Color], num_colors: int) -> Iterator[Color]:
|
|
"""Cycle colors until achieving max of `num_colors` or length of `colors`.
|
|
|
|
Extra colors will be ignored by matplotlib if there are more colors
|
|
than needed and nothing needs to be done here.
|
|
"""
|
|
max_colors = max(num_colors, len(colors))
|
|
yield from itertools.islice(itertools.cycle(colors), max_colors)
|
|
|
|
|
|
def _get_colors_from_colormap(
|
|
colormap: str | Colormap,
|
|
num_colors: int,
|
|
) -> list[Color]:
|
|
"""Get colors from colormap."""
|
|
cmap = _get_cmap_instance(colormap)
|
|
return [cmap(num) for num in np.linspace(0, 1, num=num_colors)]
|
|
|
|
|
|
def _get_cmap_instance(colormap: str | Colormap) -> Colormap:
|
|
"""Get instance of matplotlib colormap."""
|
|
if isinstance(colormap, str):
|
|
cmap = colormap
|
|
colormap = mpl.colormaps[colormap]
|
|
if colormap is None:
|
|
raise ValueError(f"Colormap {cmap} is not recognized")
|
|
return colormap
|
|
|
|
|
|
def _get_colors_from_color(
|
|
color: Color | Collection[Color],
|
|
) -> list[Color]:
|
|
"""Get colors from user input color."""
|
|
if len(color) == 0:
|
|
raise ValueError(f"Invalid color argument: {color}")
|
|
|
|
if _is_single_color(color):
|
|
color = cast(Color, color)
|
|
return [color]
|
|
|
|
color = cast(Collection[Color], color)
|
|
return list(_gen_list_of_colors_from_iterable(color))
|
|
|
|
|
|
def _is_single_color(color: Color | Collection[Color]) -> bool:
|
|
"""Check if `color` is a single color, not a sequence of colors.
|
|
|
|
Single color is of these kinds:
|
|
- Named color "red", "C0", "firebrick"
|
|
- Alias "g"
|
|
- Sequence of floats, such as (0.1, 0.2, 0.3) or (0.1, 0.2, 0.3, 0.4).
|
|
|
|
See Also
|
|
--------
|
|
_is_single_string_color
|
|
"""
|
|
if isinstance(color, str) and _is_single_string_color(color):
|
|
# GH #36972
|
|
return True
|
|
|
|
if _is_floats_color(color):
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def _gen_list_of_colors_from_iterable(color: Collection[Color]) -> Iterator[Color]:
|
|
"""
|
|
Yield colors from string of several letters or from collection of colors.
|
|
"""
|
|
for x in color:
|
|
if _is_single_color(x):
|
|
yield x
|
|
else:
|
|
raise ValueError(f"Invalid color {x}")
|
|
|
|
|
|
def _is_floats_color(color: Color | Collection[Color]) -> bool:
|
|
"""Check if color comprises a sequence of floats representing color."""
|
|
return bool(
|
|
is_list_like(color)
|
|
and (len(color) == 3 or len(color) == 4)
|
|
and all(isinstance(x, (int, float)) for x in color)
|
|
)
|
|
|
|
|
|
def _get_colors_from_color_type(color_type: str, num_colors: int) -> list[Color]:
|
|
"""Get colors from user input color type."""
|
|
if color_type == "default":
|
|
return _get_default_colors(num_colors)
|
|
elif color_type == "random":
|
|
return _get_random_colors(num_colors)
|
|
else:
|
|
raise ValueError("color_type must be either 'default' or 'random'")
|
|
|
|
|
|
def _get_default_colors(num_colors: int) -> list[Color]:
|
|
"""Get `num_colors` of default colors from matplotlib rc params."""
|
|
import matplotlib.pyplot as plt
|
|
|
|
colors = [c["color"] for c in plt.rcParams["axes.prop_cycle"]]
|
|
return colors[0:num_colors]
|
|
|
|
|
|
def _get_random_colors(num_colors: int) -> list[Color]:
|
|
"""Get `num_colors` of random colors."""
|
|
return [_random_color(num) for num in range(num_colors)]
|
|
|
|
|
|
def _random_color(column: int) -> list[float]:
|
|
"""Get a random color represented as a list of length 3"""
|
|
# GH17525 use common._random_state to avoid resetting the seed
|
|
rs = com.random_state(column)
|
|
return rs.rand(3).tolist()
|
|
|
|
|
|
def _is_single_string_color(color: Color) -> bool:
|
|
"""Check if `color` is a single string color.
|
|
|
|
Examples of single string colors:
|
|
- 'r'
|
|
- 'g'
|
|
- 'red'
|
|
- 'green'
|
|
- 'C3'
|
|
- 'firebrick'
|
|
|
|
Parameters
|
|
----------
|
|
color : Color
|
|
Color string or sequence of floats.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
True if `color` looks like a valid color.
|
|
False otherwise.
|
|
"""
|
|
conv = matplotlib.colors.ColorConverter()
|
|
try:
|
|
# error: Argument 1 to "to_rgba" of "ColorConverter" has incompatible type
|
|
# "str | Sequence[float]"; expected "tuple[float, float, float] | ..."
|
|
conv.to_rgba(color) # type: ignore[arg-type]
|
|
except ValueError:
|
|
return False
|
|
else:
|
|
return True
|