style.py 8.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278
  1. from __future__ import annotations
  2. from collections.abc import (
  3. Collection,
  4. Iterator,
  5. )
  6. import itertools
  7. from typing import (
  8. TYPE_CHECKING,
  9. cast,
  10. )
  11. import warnings
  12. import matplotlib as mpl
  13. import matplotlib.colors
  14. import numpy as np
  15. from pandas._typing import MatplotlibColor as Color
  16. from pandas.util._exceptions import find_stack_level
  17. from pandas.core.dtypes.common import is_list_like
  18. import pandas.core.common as com
  19. if TYPE_CHECKING:
  20. from matplotlib.colors import Colormap
  21. def get_standard_colors(
  22. num_colors: int,
  23. colormap: Colormap | None = None,
  24. color_type: str = "default",
  25. color: dict[str, Color] | Color | Collection[Color] | None = None,
  26. ):
  27. """
  28. Get standard colors based on `colormap`, `color_type` or `color` inputs.
  29. Parameters
  30. ----------
  31. num_colors : int
  32. Minimum number of colors to be returned.
  33. Ignored if `color` is a dictionary.
  34. colormap : :py:class:`matplotlib.colors.Colormap`, optional
  35. Matplotlib colormap.
  36. When provided, the resulting colors will be derived from the colormap.
  37. color_type : {"default", "random"}, optional
  38. Type of colors to derive. Used if provided `color` and `colormap` are None.
  39. Ignored if either `color` or `colormap` are not None.
  40. color : dict or str or sequence, optional
  41. Color(s) to be used for deriving sequence of colors.
  42. Can be either be a dictionary, or a single color (single color string,
  43. or sequence of floats representing a single color),
  44. or a sequence of colors.
  45. Returns
  46. -------
  47. dict or list
  48. Standard colors. Can either be a mapping if `color` was a dictionary,
  49. or a list of colors with a length of `num_colors` or more.
  50. Warns
  51. -----
  52. UserWarning
  53. If both `colormap` and `color` are provided.
  54. Parameter `color` will override.
  55. """
  56. if isinstance(color, dict):
  57. return color
  58. colors = _derive_colors(
  59. color=color,
  60. colormap=colormap,
  61. color_type=color_type,
  62. num_colors=num_colors,
  63. )
  64. return list(_cycle_colors(colors, num_colors=num_colors))
  65. def _derive_colors(
  66. *,
  67. color: Color | Collection[Color] | None,
  68. colormap: str | Colormap | None,
  69. color_type: str,
  70. num_colors: int,
  71. ) -> list[Color]:
  72. """
  73. Derive colors from either `colormap`, `color_type` or `color` inputs.
  74. Get a list of colors either from `colormap`, or from `color`,
  75. or from `color_type` (if both `colormap` and `color` are None).
  76. Parameters
  77. ----------
  78. color : str or sequence, optional
  79. Color(s) to be used for deriving sequence of colors.
  80. Can be either be a single color (single color string, or sequence of floats
  81. representing a single color), or a sequence of colors.
  82. colormap : :py:class:`matplotlib.colors.Colormap`, optional
  83. Matplotlib colormap.
  84. When provided, the resulting colors will be derived from the colormap.
  85. color_type : {"default", "random"}, optional
  86. Type of colors to derive. Used if provided `color` and `colormap` are None.
  87. Ignored if either `color` or `colormap`` are not None.
  88. num_colors : int
  89. Number of colors to be extracted.
  90. Returns
  91. -------
  92. list
  93. List of colors extracted.
  94. Warns
  95. -----
  96. UserWarning
  97. If both `colormap` and `color` are provided.
  98. Parameter `color` will override.
  99. """
  100. if color is None and colormap is not None:
  101. return _get_colors_from_colormap(colormap, num_colors=num_colors)
  102. elif color is not None:
  103. if colormap is not None:
  104. warnings.warn(
  105. "'color' and 'colormap' cannot be used simultaneously. Using 'color'",
  106. stacklevel=find_stack_level(),
  107. )
  108. return _get_colors_from_color(color)
  109. else:
  110. return _get_colors_from_color_type(color_type, num_colors=num_colors)
  111. def _cycle_colors(colors: list[Color], num_colors: int) -> Iterator[Color]:
  112. """Cycle colors until achieving max of `num_colors` or length of `colors`.
  113. Extra colors will be ignored by matplotlib if there are more colors
  114. than needed and nothing needs to be done here.
  115. """
  116. max_colors = max(num_colors, len(colors))
  117. yield from itertools.islice(itertools.cycle(colors), max_colors)
  118. def _get_colors_from_colormap(
  119. colormap: str | Colormap,
  120. num_colors: int,
  121. ) -> list[Color]:
  122. """Get colors from colormap."""
  123. cmap = _get_cmap_instance(colormap)
  124. return [cmap(num) for num in np.linspace(0, 1, num=num_colors)]
  125. def _get_cmap_instance(colormap: str | Colormap) -> Colormap:
  126. """Get instance of matplotlib colormap."""
  127. if isinstance(colormap, str):
  128. cmap = colormap
  129. colormap = mpl.colormaps[colormap]
  130. if colormap is None:
  131. raise ValueError(f"Colormap {cmap} is not recognized")
  132. return colormap
  133. def _get_colors_from_color(
  134. color: Color | Collection[Color],
  135. ) -> list[Color]:
  136. """Get colors from user input color."""
  137. if len(color) == 0:
  138. raise ValueError(f"Invalid color argument: {color}")
  139. if _is_single_color(color):
  140. color = cast(Color, color)
  141. return [color]
  142. color = cast(Collection[Color], color)
  143. return list(_gen_list_of_colors_from_iterable(color))
  144. def _is_single_color(color: Color | Collection[Color]) -> bool:
  145. """Check if `color` is a single color, not a sequence of colors.
  146. Single color is of these kinds:
  147. - Named color "red", "C0", "firebrick"
  148. - Alias "g"
  149. - Sequence of floats, such as (0.1, 0.2, 0.3) or (0.1, 0.2, 0.3, 0.4).
  150. See Also
  151. --------
  152. _is_single_string_color
  153. """
  154. if isinstance(color, str) and _is_single_string_color(color):
  155. # GH #36972
  156. return True
  157. if _is_floats_color(color):
  158. return True
  159. return False
  160. def _gen_list_of_colors_from_iterable(color: Collection[Color]) -> Iterator[Color]:
  161. """
  162. Yield colors from string of several letters or from collection of colors.
  163. """
  164. for x in color:
  165. if _is_single_color(x):
  166. yield x
  167. else:
  168. raise ValueError(f"Invalid color {x}")
  169. def _is_floats_color(color: Color | Collection[Color]) -> bool:
  170. """Check if color comprises a sequence of floats representing color."""
  171. return bool(
  172. is_list_like(color)
  173. and (len(color) == 3 or len(color) == 4)
  174. and all(isinstance(x, (int, float)) for x in color)
  175. )
  176. def _get_colors_from_color_type(color_type: str, num_colors: int) -> list[Color]:
  177. """Get colors from user input color type."""
  178. if color_type == "default":
  179. return _get_default_colors(num_colors)
  180. elif color_type == "random":
  181. return _get_random_colors(num_colors)
  182. else:
  183. raise ValueError("color_type must be either 'default' or 'random'")
  184. def _get_default_colors(num_colors: int) -> list[Color]:
  185. """Get `num_colors` of default colors from matplotlib rc params."""
  186. import matplotlib.pyplot as plt
  187. colors = [c["color"] for c in plt.rcParams["axes.prop_cycle"]]
  188. return colors[0:num_colors]
  189. def _get_random_colors(num_colors: int) -> list[Color]:
  190. """Get `num_colors` of random colors."""
  191. return [_random_color(num) for num in range(num_colors)]
  192. def _random_color(column: int) -> list[float]:
  193. """Get a random color represented as a list of length 3"""
  194. # GH17525 use common._random_state to avoid resetting the seed
  195. rs = com.random_state(column)
  196. return rs.rand(3).tolist()
  197. def _is_single_string_color(color: Color) -> bool:
  198. """Check if `color` is a single string color.
  199. Examples of single string colors:
  200. - 'r'
  201. - 'g'
  202. - 'red'
  203. - 'green'
  204. - 'C3'
  205. - 'firebrick'
  206. Parameters
  207. ----------
  208. color : Color
  209. Color string or sequence of floats.
  210. Returns
  211. -------
  212. bool
  213. True if `color` looks like a valid color.
  214. False otherwise.
  215. """
  216. conv = matplotlib.colors.ColorConverter()
  217. try:
  218. # error: Argument 1 to "to_rgba" of "ColorConverter" has incompatible type
  219. # "str | Sequence[float]"; expected "tuple[float, float, float] | ..."
  220. conv.to_rgba(color) # type: ignore[arg-type]
  221. except ValueError:
  222. return False
  223. else:
  224. return True