Coverage for E:\1vsCode\python\printbuddies\src\printbuddies\gradient.py: 42%

93 statements  

« prev     ^ index     » next       coverage.py v7.2.2, created at 2024-02-15 19:42 -0600

1import itertools 

2import string 

3from dataclasses import dataclass 

4from typing import Sequence, SupportsIndex 

5 

6from rich.color import Color 

7from typing_extensions import Self 

8 

9from .colormap import Tag 

10 

11 

12@dataclass 

13class RGB: 

14 """ 

15 Dataclass representing a 3 channel RGB color that is converted to a `rich` tag when casted to a string. 

16 

17 >>> color = RGB(100, 100, 100) 

18 >>> str(color) 

19 >>> "[rgb(100,100,100)]" 

20 >>> from rich.console import Console 

21 >>> console = Console() 

22 >>> console.print(f"{color}Yeehaw") 

23 

24 Can also be initialized using a color name from https://rich.readthedocs.io/en/stable/appendix/colors.html 

25 

26 >>> color = RGB(name="magenta3") 

27 >>> print(color) 

28 >>> "[rgb(215,0,215)]" 

29 

30 Supports addition and subtraction of `RGB` objects as well as scalar multiplication and division. 

31 

32 >>> color1 = RGB(100, 100, 100) 

33 >>> color2 = RGB(25, 50, 75) 

34 >>> print(color1 + color2) 

35 >>> "[rgb(125,150,175)]" 

36 >>> print(color2 * 2) 

37 >>> "[rgb(50,100,150)]" 

38 """ 

39 

40 # Typing these as floats so `Gradient` can fractionally increment them 

41 # When casted to a string, the values will be rounded to integers 

42 r: float = 0 

43 g: float = 0 

44 b: float = 0 

45 name: str = "" 

46 

47 def __post_init__(self): 

48 if self.name: 

49 self.r, self.g, self.b = Color.parse(self.name).get_truecolor() 

50 

51 def __str__(self) -> str: 

52 return f"[rgb({round(self.r)},{round(self.g)},{round(self.b)})]" 

53 

54 def __sub__(self, other: Self) -> Self: 

55 return self.__class__(self.r - other.r, self.g - other.g, self.b - other.b) 

56 

57 def __add__(self, other: Self) -> Self: 

58 return self.__class__(self.r + other.r, self.g + other.g, self.b + other.b) 

59 

60 def __truediv__(self, val: float) -> Self: 

61 return self.__class__(self.r / val, self.g / val, self.b / val) 

62 

63 def __mul__(self, val: float) -> Self: 

64 return self.__class__(self.r * val, self.g * val, self.b * val) 

65 

66 def __eq__(self, other: Self) -> bool: 

67 return all(getattr(self, c) == getattr(other, c) for c in "rgb") 

68 

69 

70ColorType = RGB | tuple[int, int, int] | str | Tag 

71 

72 

73class _Blender: 

74 """ 

75 Apply a color blend from a start color to a stop color across text when printed with the `rich` package. 

76 """ 

77 

78 def __init__( 

79 self, 

80 start: RGB, 

81 stop: RGB, 

82 ): 

83 self.start = start 

84 self.stop = stop 

85 

86 @property 

87 def valid_characters(self) -> str: 

88 """Characters a color step can be applied to.""" 

89 return string.ascii_letters + string.digits + string.punctuation 

90 

91 def _get_step_sizes(self, num_steps: int) -> RGB: 

92 """Returns a `RGB` object representing the step size for each color channel.""" 

93 return (self.stop - self.start) / num_steps 

94 

95 def _get_blended_color(self, step: int, step_sizes: RGB) -> RGB: 

96 """Returns a `RGB` object representing the color at `step`.""" 

97 return self.start + (step_sizes * step) 

98 

99 def _get_num_steps(self, text: str) -> int: 

100 """Returns the number of steps the blend should be divided into.""" 

101 return len([ch for ch in text if ch in self.valid_characters]) - 1 

102 

103 def apply(self, text: str) -> str: 

104 """Apply the blend to ascii letters, digits, and punctuation in `text`.""" 

105 num_steps = self._get_num_steps(text) 

106 if num_steps < 0: # no valid characters 

107 return text 

108 elif num_steps == 0: # one valid character, just apply start color 

109 return f"{self.start}{text}[/]" 

110 step_sizes = self._get_step_sizes(num_steps) 

111 blended_text = "" 

112 step = 0 

113 for ch in text: 

114 if ch in self.valid_characters: 

115 blended_text += f"{self._get_blended_color(step, step_sizes)}{ch}[/]" 

116 step += 1 

117 else: 

118 blended_text += ch 

119 return blended_text 

120 

121 

122class Gradient(list[RGB]): 

123 """ 

124 Apply an arbitrary number of color gradients to strings when using `rich`. 

125 

126 When applied to a string, each character will increment in color from a start to a stop color. 

127 

128 Colors can be specified by either 

129 a 3 tuple representing RGB values, 

130 a `pocketchange.RGB` object, 

131 a `pocketchange.Tag` object, 

132 or a color name from https://rich.readthedocs.io/en/stable/appendix/colors.html. 

133 

134 Tuple: 

135 >>> gradient = Gradient([(255, 0, 0), (0, 255, 0)]) 

136 

137 `pocketchange.RGB`: 

138 >>> gradient = Gradient([RGB(255, 0, 0), RGB(0, 255, 0)]) 

139 

140 `pocketchange.Tag`: 

141 >>> colors = pocketchange.ColorMap() 

142 >>> gradient = Gradient([colors.red, colors.green]) 

143 

144 Name: 

145 >>> gradient = Gradient(["red", "green"]) 

146 

147 Usage: 

148 >>> from pocketchange import Gradient 

149 >>> from rich.console import Console 

150 >>> 

151 >>> console = Console() 

152 >>> gradient = Gradient(["red", "green"]) 

153 >>> text = "Yeehaw" 

154 >>> gradient_text = gradient.apply(text) 

155 >>> # This produces: 

156 >>> print(gradient_text) 

157 >>> "[rgb(128,0,0)]Y[/][rgb(102,25,0)]e[/][rgb(76,51,0)]e[/][rgb(51,76,0)]h[/][rgb(25,102,0)]a[/][rgb(0,128,0)]w[/]" 

158 >>> 

159 >>> # When used with `console.print`, each character will be a different color 

160 >>> console.print(gradient_text) 

161 >>> 

162 >>> # `Gradient` inherits from `list` so colors may be appended, inserted, or extended 

163 >>> gradient.append("blue") 

164 >>> print(gradient.apply(text)) 

165 >>> "[rgb(128,0,0)]Y[/][rgb(64,64,0)]e[/][rgb(0,128,0)]e[/][rgb(0,128,0)]h[/][rgb(0,64,64)]a[/][rgb(0,0,128)]w[/]" 

166 >>> print(gradient) 

167 >>> [RGB(r=128, g=0, b=0, name='red'), RGB(r=0, g=128, b=0, name='green'), RGB(r=0, g=0, b=128, name='blue')] 

168 >>> 

169 >>> Gradient(gradient + gradient[1::-1]) 

170 >>> [RGB(r=128, g=0, b=0, name='red'), RGB(r=0, g=128, b=0, name='green'), RGB(r=0, g=0, b=128, name='blue'), RGB(r=0, g=128, b=0, name='green'), RGB(r=128, g=0, b=0, name='red')] 

171 

172 """ 

173 

174 def __init__(self, colors: Sequence[ColorType] = ["pink1", "turquoise2"]): 

175 colors_ = [self._parse(color) for color in colors] 

176 super().__init__(colors_) 

177 

178 def _parse(self, color: ColorType) -> RGB: 

179 if isinstance(color, RGB): 

180 return color 

181 elif isinstance(color, str): 

182 return RGB(name=color) 

183 elif isinstance(color, Tag): 

184 return RGB(name=color.name) 

185 elif isinstance(color, tuple): 

186 return RGB(*color) 

187 raise ValueError(f"{color!r} is an invalid type.") 

188 

189 def __setitem__(self, index: int, color: ColorType): 

190 super().__setitem__(index, self._parse(color)) 

191 

192 def append(self, color: ColorType): 

193 super().append(self._parse(color)) 

194 

195 def insert(self, index: SupportsIndex, color: ColorType): 

196 super().insert(index, self._parse(color)) 

197 

198 def extend(self, colors: list[ColorType]): 

199 super().extend([self._parse(color) for color in colors]) 

200 

201 def _get_blenders(self) -> list[_Blender]: 

202 return [_Blender(colors[0], colors[1]) for colors in itertools.pairwise(self)] 

203 

204 def _batch_text(self, text: str, n: int) -> list[str]: 

205 """Split `text` into `n` chunks. 

206 

207 All chunks will be the same size, except potentially the last chunk.""" 

208 batch_size = int(len(text) / n) 

209 if batch_size == 0: 

210 return [ch for ch in text] 

211 batched_text = [ 

212 text[i * batch_size : (i * batch_size) + batch_size] for i in range(n - 1) 

213 ] 

214 lastdex = n - 1 

215 batched_text.append(text[lastdex * batch_size :]) 

216 return batched_text 

217 

218 def apply(self, text: str) -> str: 

219 """Format `text` such that, when printed with `rich`, 

220 the displayed text changes colors according to this instance's color list.""" 

221 blenders = self._get_blenders() 

222 batches = self._batch_text(text, len(blenders)) 

223 return "".join( 

224 blender.apply(batch) for blender, batch in zip(blenders, batches) 

225 )