Coverage for src\shortrich\gradient.py: 95%

59 statements  

« prev     ^ index     » next       coverage.py v7.2.2, created at 2024-02-13 11:18 -0600

1from dataclasses import dataclass 

2 

3from rich.color import Color 

4from typing_extensions import Self 

5 

6from .colormap import Tag 

7 

8 

9@dataclass 

10class RGB: 

11 # typing these as floats so `Gradient` can fractionally increment them 

12 # but when casted to a string, the values will be converted to ints 

13 r: float = 0 

14 g: float = 0 

15 b: float = 0 

16 name: str = "" 

17 

18 def __post_init__(self): 

19 if self.name: 

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

21 

22 def __str__(self) -> str: 

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

24 

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

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

27 

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

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

30 

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

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

33 

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

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

36 

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

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

39 

40 

41class Gradient: 

42 def __init__( 

43 self, 

44 start: tuple[int, int, int] | str | Tag, 

45 stop: tuple[int, int, int] | str | Tag, 

46 ): 

47 self._start = self._parse(start) 

48 self._stop = self._parse(stop) 

49 

50 @property 

51 def start(self) -> RGB: 

52 return self._start 

53 

54 @start.setter 

55 def start(self, color: str | Tag | tuple[int, int, int]): 

56 self._start = self._parse(color) 

57 

58 @property 

59 def stop(self) -> RGB: 

60 return self._stop 

61 

62 @stop.setter 

63 def stop(self, color: str | Tag | tuple[int, int, int]): 

64 self._stop = self._parse(color) 

65 

66 def _parse(self, color: str | Tag | tuple[int, int, int]) -> RGB: 

67 if isinstance(color, str): 

68 return RGB(name=color) 

69 elif isinstance(color, Tag): 

70 return RGB(name=color.name) 

71 elif isinstance(color, tuple): 

72 return RGB(*color) 

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

74 

75 def _get_step_sizes(self, total_steps: int) -> RGB: 

76 return (self.stop - self.start) / total_steps 

77 

78 def _get_gradient_color(self, step: int, step_sizes: RGB) -> RGB: 

79 return self.start + (step_sizes * step) 

80 

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

82 return len(text) - 1 

83 

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

85 steps = self._get_num_steps(text) 

86 step_sizes = self._get_step_sizes(steps) 

87 return "".join( 

88 f"{self._get_gradient_color(i, step_sizes)}{ch}[/]" 

89 for i, ch in enumerate(text) 

90 )