color_util.py 2.9 KB

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