Coverage for app/chart.py: 50%

1import colorsys 

2import random 

3from datetime import datetime 

4from zoneinfo import ZoneInfo 

5 

6import altair as alt 

7import pandas as pd 

8from sqlalchemy import text 

9 

10from app.core.db import read_engine 

11from app.core.logger import logger 

12from app.models import QueryParams 

13from app.utils import start_of_week 

14 

15 

16def generate_altair_colors(n: int, seed: int = 42) -> list[str]: 

17 random.seed(seed) 

18 colors: list[str] = [] 

19 

20 while len(colors) < n: 

21 # Python blue, yellow, green 

22 if len(colors) == 0: 

23 colors.append("#346f9e") 

24 continue 

25 elif len(colors) == 1: 

26 colors.append("#ffde56") 

27 continue 

28 elif len(colors) == 2: 

29 colors.append("#2e7d32") 

30 continue 

31 

32 h = random.uniform(0, 1) 

33 s = random.uniform(0.5, 0.9) 

34 v = random.uniform(0.6, 0.9) 

35 rgb = colorsys.hsv_to_rgb(h, s, v) 

36 hex_color = "#{:02x}{:02x}{:02x}".format( 

37 int(rgb[0] * 255), int(rgb[1] * 255), int(rgb[2] * 255) 

38 ) 

39 if hex_color not in colors: 

40 colors.append(hex_color) 

41 

42 return colors 

43 

44 

45def generate_chart(query_params: QueryParams, theme: str) -> alt.Chart: 

46 start_week = start_of_week(query_params.time_range.date_str) 

47 end_week = start_of_week(datetime.now(ZoneInfo("UTC"))) 

48 logger.info( 

49 f"Generating chart for packages: {query_params.packages}, " 

50 f"time range: {query_params.time_range.value}, theme: {theme}, " 

51 f"date range: {start_week} to {end_week}" 

52 ) 

53 packages = {str(i): package for i, package in enumerate(query_params.packages)} 

54 placeholders = ", ".join(f":{i}" for i in range(len(query_params.packages))) 

55 theme_config = { 

56 "light": {"theme": "default", "tooltip": "dark"}, 

57 "dark": {"theme": "dark", "tooltip": "white"}, 

58 } 

59 alt.theme.enable(theme_config[theme]["theme"]) 

60 alt.ViewBackground(fillOpacity=0) 

61 

62 with read_engine.connect() as conn: 

63 result = conn.execute( 

64 text(f""" 

65 select 

66 package_name, 

67 package_downloaded_week, 

68 downloads, 

69 downloads * 1.0 / nullif(sum(downloads) over(partition by package_downloaded_week), 0) as percentage_of_total_downloads, 

70 cumulative_downloads, 

71 weeks_since_first_distribution 

72 from 

73 pypi_package_downloads_weekly_metrics 

74 where true 

75 and package_name in ({placeholders}) 

76 and package_downloaded_week >= :start_week 

77 and package_downloaded_week < :end_week 

78 """), 

79 {**packages, "start_week": start_week, "end_week": end_week}, 

80 ) 

81 downloads = result.fetchall() 

82 

83 df = pd.DataFrame( 

84 downloads, 

85 columns=[ 

86 "package", 

87 "week", 

88 "downloads", 

89 "percentage_of_total_downloads", 

90 "cumulative_downloads", 

91 "weeks_since_first_distribution", 

92 ], 

93 ) 

94 df["week"] = pd.to_datetime(df["week"]) 

95 

96 highlight = alt.selection_point(on="pointerover", fields=["package"], nearest=True) 

97 

98 x_weeks_since_first_release = dict( 

99 shorthand="weeks_since_first_distribution:Q", 

100 title="Weeks Since First Release", 

101 format=",", 

102 ) 

103 

104 x_week = dict( 

105 shorthand="week:T", 

106 title="", 

107 format="%Y-%m-%d", 

108 ) 

109 

110 if query_params.time_range.value == "allTimeCumulativeAlignTimeline": 

111 x = x_weeks_since_first_release 

112 else: 

113 x = x_week 

114 

115 y_cumulative_downloads = dict( 

116 shorthand="cumulative_downloads:Q", title="Cumulative Downloads", format="," 

117 ) 

118 

119 y_perecent_of_total = dict( 

120 shorthand="percentage_of_total_downloads:Q", 

121 title="Percent of Total Downloads", 

122 format=".2%", 

123 ) 

124 

125 y_weekly_downloads = dict(shorthand="downloads:Q", title="Downloads", format=",") 

126 

127 is_cumulative = query_params.time_range.value in ( 

128 "allTimeCumulative", 

129 "allTimeCumulativeAlignTimeline", 

130 ) 

131 show_percentage = query_params.show_percentage == "on" and not is_cumulative 

132 

133 if is_cumulative: 

134 y = y_cumulative_downloads 

135 elif show_percentage: 

136 y = y_perecent_of_total 

137 else: 

138 y = y_weekly_downloads 

139 

140 base = alt.Chart(df).encode( 

141 x=alt.X( 

142 x["shorthand"], 

143 title=x["title"], 

144 axis=alt.Axis(tickCount=2.5, labelFontSize=14), 

145 ), 

146 y=alt.Y( 

147 y["shorthand"], 

148 title=y["title"], 

149 axis=alt.Axis( 

150 tickCount=3, 

151 labelFontSize=14, 

152 format=".0%" if show_percentage else alt.Undefined, 

153 ), 

154 scale=alt.Scale(domain=[0, 1]) if show_percentage else alt.Undefined, 

155 ), 

156 color="package:N", 

157 tooltip=[ 

158 alt.Tooltip("package:N"), 

159 alt.Tooltip(**x_week), 

160 alt.Tooltip(**y_weekly_downloads), 

161 alt.Tooltip(**y_cumulative_downloads), 

162 alt.Tooltip(**y_perecent_of_total), 

163 alt.Tooltip(**x_weeks_since_first_release), 

164 ], 

165 ) 

166 

167 points = base.mark_circle().encode(opacity=alt.value(0)).add_params(highlight) 

168 

169 lines = base.mark_line( 

170 interpolate="monotone", strokeWidth=4, strokeCap="round" 

171 ).encode( 

172 size=alt.when(~highlight).then(alt.value(2)).otherwise(alt.value(3)), 

173 color=alt.Color( 

174 "package:N", 

175 scale=alt.Scale( 

176 domain=packages.values(), 

177 range=generate_altair_colors(len(packages)), 

178 ), 

179 ).legend(orient="top", title=None, labelFontSize=14), 

180 ) 

181 

182 return ( 

183 (points + lines) 

184 .properties( 

185 width="container", 

186 height=400, 

187 usermeta={ 

188 "embedOptions": { 

189 "tooltip": {"theme": theme_config[theme]["tooltip"]}, 

190 "actions": False, 

191 }, 

192 }, 

193 ) 

194 .configure(background="rgba(0,0,0,0)") 

195 )