Coverage for apps/character.py: 100%

1"""Character information for StreamWise.""" 

2 

3import logging 

4import random 

5 

6from typing import Dict 

7from typing import Optional 

8from typing import Any 

9 

10from PIL import Image 

11 

12# https://github.com/hexgrad/kokoro/blob/main/kokoro.js/src/voices.js 

13# https://github.com/hexgrad/kokoro/blob/main/kokoro.js/README.md 

14KOKORO_VOICES = { 

15 "female": [ 

16 "af_heart", 

17 "af_alloy", 

18 "af_aoede", 

19 "af_bella", 

20 "af_jessica", 

21 "af_kore", 

22 "af_nicole", 

23 "af_nova", 

24 "af_river", 

25 "af_sarah", 

26 "af_sky" 

27 ], 

28 "male": [ 

29 "am_adam", 

30 "am_echo", 

31 "am_eric", 

32 "am_fenrir", 

33 "am_liam", 

34 "am_michael", 

35 "am_onyx", 

36 "am_puck", 

37 "am_santa" 

38 ] 

39} 

40 

41VIBEVOICE_VOICES = { 

42 "Female": [ 

43 "af_bella", 

44 "af_heart", 

45 "af_kore", 

46 "af_nicole", 

47 "woman_000", 

48 ], 

49 "Male": [ 

50 "am_adam", 

51 "am_fenrir", 

52 "am_michael", 

53 "am_puck", 

54 "girish_001", 

55 "girish_002", 

56 "ricardo_001", 

57 "ricardo_002", 

58 ] 

59} 

60 

61 

62# Comes from javascript/bootstrap: 

63# const colorClasses = ["text-primary", "text-success", "text-danger", "text-warning", "text-info", "text-secondary"]; 

64CHARACTER_COLOR_CLASSES = { 

65 "text-primary": "#0d6efd", # blue 

66 "text-success": "#198754", # green 

67 "text-danger": "#dc3545", # red 

68 "text-warning": "#ffc107", # yellow 

69 "text-info": "#0dcaf0", # cyan 

70 "text-secondary": "#6c757d", # gray 

71} 

72 

73 

74class Character: 

75 """Character information.""" 

76 def __init__( 

77 self, 

78 name: str, 

79 gender: str = "Unknown", 

80 description: Optional[str] = None, 

81 speech_speed: float = 1.1, 

82 ) -> None: 

83 self.name = name 

84 self.gender = gender 

85 self.voice = self.get_random_voice() 

86 self.speech_speed = speech_speed 

87 self.image: Optional[Image.Image] = None 

88 self.description = description 

89 

90 def get_random_voice(self) -> str: 

91 """Get a random voice based on the gender.""" 

92 gender = self.gender.lower() 

93 gender_voices = KOKORO_VOICES.get(gender, []) 

94 if not gender_voices: 

95 logging.warning(f"Unknown gender: {self.gender}.") 

96 return "af_heart" 

97 return random.choice(gender_voices) 

98 

99 

100class Characters: 

101 """Manage multiple characters.""" 

102 def __init__(self) -> None: 

103 self.characters: Dict[str, Character] = {} 

104 

105 def __len__(self) -> int: 

106 return len(self.characters) 

107 

108 def __contains__(self, name: str) -> bool: 

109 return name in self.characters 

110 

111 def __getitem__(self, name: str) -> Character: 

112 if name not in self.characters: 

113 raise KeyError(f"Character '{name}' not found.") 

114 return self.characters[name] 

115 

116 def __setitem__(self, name: str, character: Character) -> None: 

117 if not isinstance(character, Character): 

118 raise ValueError("character must be an instance of Character") 

119 self.characters[name] = character 

120 

121 def __iter__(self) -> Any: 

122 return iter(self.characters.values()) 

123 

124 def __repr__(self) -> str: 

125 return f"Characters({list(self.characters.keys())})" 

126 

127 def get_index(self, name: str) -> int: 

128 """Get index of character by name.""" 

129 if name not in self.characters: 

130 return 0 

131 return list(self.characters).index(name) 

132 

133 def get_by_index(self, index: int) -> Character: 

134 """Get character by index.""" 

135 if index < 0 or index >= len(self.characters): 

136 raise IndexError("Character index out of range") 

137 character = list(self.characters.values())[index] 

138 return character 

139 

140 def get_position( 

141 self, 

142 name: str, 

143 ) -> str: 

144 """Get character position in the image: left, center, right.""" 

145 index = self.get_index(name) 

146 num_characters = len(self.characters) 

147 

148 if num_characters == 1: 

149 return "center" 

150 if num_characters == 2: 

151 if index == 0: 

152 return "left" 

153 elif index == 1: 

154 return "right" 

155 if num_characters == 3: 

156 if index == 0: 

157 return "left" 

158 elif index == 1: 

159 return "center" 

160 elif index == 2: 

161 return "right" 

162 return "left" 

163 

164 def get_color(self, name: str) -> str: 

165 index = self.get_index(name) % len(CHARACTER_COLOR_CLASSES) 

166 color = CHARACTER_COLOR_CLASSES[list(CHARACTER_COLOR_CLASSES.keys())[index]] 

167 return color