palette.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. from math import sqrt
  2. from functools import lru_cache
  3. from typing import Sequence, Tuple, TYPE_CHECKING
  4. from .color_triplet import ColorTriplet
  5. if TYPE_CHECKING:
  6. from pip._vendor.rich.table import Table
  7. class Palette:
  8. """A palette of available colors."""
  9. def __init__(self, colors: Sequence[Tuple[int, int, int]]):
  10. self._colors = colors
  11. def __getitem__(self, number: int) -> ColorTriplet:
  12. return ColorTriplet(*self._colors[number])
  13. def __rich__(self) -> "Table":
  14. from pip._vendor.rich.color import Color
  15. from pip._vendor.rich.style import Style
  16. from pip._vendor.rich.text import Text
  17. from pip._vendor.rich.table import Table
  18. table = Table(
  19. "index",
  20. "RGB",
  21. "Color",
  22. title="Palette",
  23. caption=f"{len(self._colors)} colors",
  24. highlight=True,
  25. caption_justify="right",
  26. )
  27. for index, color in enumerate(self._colors):
  28. table.add_row(
  29. str(index),
  30. repr(color),
  31. Text(" " * 16, style=Style(bgcolor=Color.from_rgb(*color))),
  32. )
  33. return table
  34. # This is somewhat inefficient and needs caching
  35. @lru_cache(maxsize=1024)
  36. def match(self, color: Tuple[int, int, int]) -> int:
  37. """Find a color from a palette that most closely matches a given color.
  38. Args:
  39. color (Tuple[int, int, int]): RGB components in range 0 > 255.
  40. Returns:
  41. int: Index of closes matching color.
  42. """
  43. red1, green1, blue1 = color
  44. _sqrt = sqrt
  45. get_color = self._colors.__getitem__
  46. def get_color_distance(index: int) -> float:
  47. """Get the distance to a color."""
  48. red2, green2, blue2 = get_color(index)
  49. red_mean = (red1 + red2) // 2
  50. red = red1 - red2
  51. green = green1 - green2
  52. blue = blue1 - blue2
  53. return _sqrt(
  54. (((512 + red_mean) * red * red) >> 8)
  55. + 4 * green * green
  56. + (((767 - red_mean) * blue * blue) >> 8)
  57. )
  58. min_index = min(range(len(self._colors)), key=get_color_distance)
  59. return min_index
  60. if __name__ == "__main__": # pragma: no cover
  61. import colorsys
  62. from typing import Iterable
  63. from pip._vendor.rich.color import Color
  64. from pip._vendor.rich.console import Console, ConsoleOptions
  65. from pip._vendor.rich.segment import Segment
  66. from pip._vendor.rich.style import Style
  67. class ColorBox:
  68. def __rich_console__(
  69. self, console: Console, options: ConsoleOptions
  70. ) -> Iterable[Segment]:
  71. height = console.size.height - 3
  72. for y in range(0, height):
  73. for x in range(options.max_width):
  74. h = x / options.max_width
  75. l = y / (height + 1)
  76. r1, g1, b1 = colorsys.hls_to_rgb(h, l, 1.0)
  77. r2, g2, b2 = colorsys.hls_to_rgb(h, l + (1 / height / 2), 1.0)
  78. bgcolor = Color.from_rgb(r1 * 255, g1 * 255, b1 * 255)
  79. color = Color.from_rgb(r2 * 255, g2 * 255, b2 * 255)
  80. yield Segment("▄", Style(color=color, bgcolor=bgcolor))
  81. yield Segment.line()
  82. console = Console()
  83. console.print(ColorBox())