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
« 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
6from rich.color import Color
7from typing_extensions import Self
9from .colormap import Tag
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.
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")
24 Can also be initialized using a color name from https://rich.readthedocs.io/en/stable/appendix/colors.html
26 >>> color = RGB(name="magenta3")
27 >>> print(color)
28 >>> "[rgb(215,0,215)]"
30 Supports addition and subtraction of `RGB` objects as well as scalar multiplication and division.
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 """
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 = ""
47 def __post_init__(self):
48 if self.name:
49 self.r, self.g, self.b = Color.parse(self.name).get_truecolor()
51 def __str__(self) -> str:
52 return f"[rgb({round(self.r)},{round(self.g)},{round(self.b)})]"
54 def __sub__(self, other: Self) -> Self:
55 return self.__class__(self.r - other.r, self.g - other.g, self.b - other.b)
57 def __add__(self, other: Self) -> Self:
58 return self.__class__(self.r + other.r, self.g + other.g, self.b + other.b)
60 def __truediv__(self, val: float) -> Self:
61 return self.__class__(self.r / val, self.g / val, self.b / val)
63 def __mul__(self, val: float) -> Self:
64 return self.__class__(self.r * val, self.g * val, self.b * val)
66 def __eq__(self, other: Self) -> bool:
67 return all(getattr(self, c) == getattr(other, c) for c in "rgb")
70ColorType = RGB | tuple[int, int, int] | str | Tag
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 """
78 def __init__(
79 self,
80 start: RGB,
81 stop: RGB,
82 ):
83 self.start = start
84 self.stop = stop
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
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
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)
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
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
122class Gradient(list[RGB]):
123 """
124 Apply an arbitrary number of color gradients to strings when using `rich`.
126 When applied to a string, each character will increment in color from a start to a stop color.
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.
134 Tuple:
135 >>> gradient = Gradient([(255, 0, 0), (0, 255, 0)])
137 `pocketchange.RGB`:
138 >>> gradient = Gradient([RGB(255, 0, 0), RGB(0, 255, 0)])
140 `pocketchange.Tag`:
141 >>> colors = pocketchange.ColorMap()
142 >>> gradient = Gradient([colors.red, colors.green])
144 Name:
145 >>> gradient = Gradient(["red", "green"])
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')]
172 """
174 def __init__(self, colors: Sequence[ColorType] = ["pink1", "turquoise2"]):
175 colors_ = [self._parse(color) for color in colors]
176 super().__init__(colors_)
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.")
189 def __setitem__(self, index: int, color: ColorType):
190 super().__setitem__(index, self._parse(color))
192 def append(self, color: ColorType):
193 super().append(self._parse(color))
195 def insert(self, index: SupportsIndex, color: ColorType):
196 super().insert(index, self._parse(color))
198 def extend(self, colors: list[ColorType]):
199 super().extend([self._parse(color) for color in colors])
201 def _get_blenders(self) -> list[_Blender]:
202 return [_Blender(colors[0], colors[1]) for colors in itertools.pairwise(self)]
204 def _batch_text(self, text: str, n: int) -> list[str]:
205 """Split `text` into `n` chunks.
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
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 )