color_util.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. from __future__ import annotations
  2. import colorsys
  3. from typing import NamedTuple
  4. from typing_extensions import Literal
  5. AnsiMode = Literal['default', 'ansi', '8bit', 'rgb']
  6. def redistribute_rgb(r: int, g: int, b: int) -> tuple[int, int, int]:
  7. """
  8. Redistribute RGB after lightening
  9. Credit: https://stackoverflow.com/a/141943/7346633
  10. """
  11. threshold = 255.999
  12. m = max(r, g, b)
  13. if m <= threshold:
  14. return int(r), int(g), int(b)
  15. total = r + g + b
  16. if total >= 3 * threshold:
  17. return int(threshold), int(threshold), int(threshold)
  18. x = (3 * threshold - total) / (3 * m - total)
  19. gray = threshold - x * m
  20. return int(gray + x * r), int(gray + x * g), int(gray + x * b)
  21. class RGB(NamedTuple):
  22. r: int
  23. g: int
  24. b: int
  25. @classmethod
  26. def from_hex(cls, hex: str) -> "RGB":
  27. """
  28. Create color from hex code
  29. >>> RGB.from_hex('#FFAAB7')
  30. RGB(r=255, g=170, b=183)
  31. :param hex: Hex color code
  32. :return: RGB object
  33. """
  34. while hex.startswith('#'):
  35. hex = hex[1:]
  36. r = int(hex[0:2], 16)
  37. g = int(hex[2:4], 16)
  38. b = int(hex[4:6], 16)
  39. return cls(r, g, b)
  40. def to_ansi_rgb(self, foreground: bool = True) -> str:
  41. """
  42. Convert RGB to ANSI TrueColor (RGB) Escape Code.
  43. This uses the 24-bit color encoding (an uint8 for each color value), and supports 16 million
  44. colors. However, not all terminal emulators support this escape code. (For example, IntelliJ
  45. debug console doesn't support it).
  46. Currently, we do not know how to detect whether a terminal environment supports ANSI RGB. If
  47. you have any thoughts, feel free to submit an issue on our Github page!
  48. :param foreground: Whether the color is for foreground text or background color
  49. :return: ANSI RGB escape code like \033[38;2;255;100;0m
  50. """
  51. c = '38' if foreground else '48'
  52. return f'\033[{c};2;{self.r};{self.g};{self.b}m'
  53. def to_ansi_8bit(self, foreground: bool = True) -> str:
  54. """
  55. Convert RGB to ANSI 8bit 256 Color Escape Code.
  56. This encoding supports 256 colors in total.
  57. :return: ANSI 256 escape code like \033[38;5;206m'
  58. """
  59. r, g, b = self.r, self.g, self.b
  60. sep = 42.5
  61. while True:
  62. if r < sep or g < sep or b < sep:
  63. gray = r < sep and g < sep and b < sep
  64. break
  65. sep += 42.5
  66. if gray:
  67. color = 232 + (r + g + b) / 33
  68. else:
  69. color = 16 + int(r / 256. * 6) * 36 + int(g / 256. * 6) * 6 + int(b / 256. * 6)
  70. c = '38' if foreground else '48'
  71. return f'\033[{c};5;{int(color)}m'
  72. def to_ansi_16(self, foreground: bool = True) -> str:
  73. """
  74. Convert RGB to ANSI 16 Color Escape Code
  75. :return: ANSI 16 escape code
  76. """
  77. raise NotImplementedError()
  78. def to_ansi(self, mode: AnsiMode, foreground: bool = True):
  79. if mode == 'rgb':
  80. return self.to_ansi_rgb(foreground)
  81. if mode == '8bit':
  82. return self.to_ansi_8bit(foreground)
  83. if mode == 'ansi':
  84. return self.to_ansi_16(foreground)
  85. def lighten(self, multiplier: float) -> 'RGB':
  86. """
  87. Lighten the color by a multiplier
  88. :param multiplier: Multiplier
  89. :return: Lightened color (original isn't modified)
  90. """
  91. return RGB(*redistribute_rgb(*[v * multiplier for v in self]))
  92. def set_light(self, light: int) -> 'RGB':
  93. """
  94. Set HSL lightness value
  95. :param light: Lightness value
  96. :return: New color (original isn't modified)
  97. """
  98. h, l, s = colorsys.rgb_to_hls(*[v / 255.0 for v in self])
  99. return RGB(*[round(v * 255.0) for v in colorsys.hls_to_rgb(h, light, s)])