Coverage for mindsdb / integrations / handlers / pypi_handler / pypi_tables.py: 0%

1from typing import List 

2 

3import pandas as pd 

4from mindsdb_sql_parser import ast 

5 

6from mindsdb.integrations.utilities.handlers.query_utilities import ( 

7 SELECTQueryExecutor, 

8 SELECTQueryParser, 

9) 

10from mindsdb.integrations.libs.api_handler import APIHandler, APITable 

11from mindsdb.integrations.utilities.sql_utils import conditions_to_filter 

12 

13 

14class CustomAPITable(APITable): 

15 name: str = None 

16 columns: List[str] = [ 

17 "category", 

18 "date", 

19 "downloads", 

20 ] 

21 

22 def __init__(self, handler: APIHandler): 

23 super().__init__(handler) 

24 self.handler.connect() 

25 

26 def get_columns(self, ignore: List[str] = []) -> List[str]: 

27 """columns 

28 

29 Args: 

30 ignore (List[str], optional): exclusion items. Defaults to []. 

31 

32 Returns: 

33 List[str]: available columns with `ignore` items removed from the list. 

34 """ 

35 

36 return [item for item in self.columns if item not in ignore] 

37 

38 

39class PyPIRecentTable(CustomAPITable): 

40 name: str = "recent" 

41 columns: List[str] = [ 

42 "last_day", 

43 "last_week", 

44 "last_month", 

45 ] 

46 

47 def select(self, query: ast.Select) -> pd.DataFrame: 

48 """triggered at the SELECT query 

49 

50 Args: 

51 query (ast.Select): user's entered query 

52 

53 Returns: 

54 pd.DataFrame: the queried information 

55 """ 

56 params = conditions_to_filter(query.where) 

57 

58 package_name = params["package"] 

59 period = params.get("period", None) 

60 all_cols = { 

61 "day": "last_day", 

62 "week": "last_week", 

63 "month": "last_month", 

64 } 

65 

66 to_be_excluded = [] 

67 

68 if period: 

69 if period in all_cols.keys(): 

70 del all_cols[period] 

71 to_be_excluded = list(all_cols.values()) 

72 else: 

73 raise ValueError( 

74 "Make sure that one of `day`, `week` or `month` values is assigned to `period`." 

75 ) 

76 

77 select_statement_parser = SELECTQueryParser( 

78 query, PyPIRecentTable.name, self.get_columns(to_be_excluded) 

79 ) 

80 ( 

81 selected_columns, 

82 _, 

83 order_by_conditions, 

84 _, 

85 ) = select_statement_parser.parse_query() 

86 

87 raw_df = self.handler.connection(name=package_name).recent(period) 

88 

89 select_statement_executor = SELECTQueryExecutor( 

90 raw_df, selected_columns, [], order_by_conditions 

91 ) 

92 

93 result_df = select_statement_executor.execute_query() 

94 

95 return result_df 

96 

97 

98class PyPIOverallTable(CustomAPITable): 

99 name: str = "overall" 

100 

101 def select(self, query: ast.Select) -> pd.DataFrame: 

102 """triggered at the SELECT query 

103 

104 Args: 

105 query (ast.Select): user's entered query 

106 

107 Returns: 

108 pd.DataFrame: the queried information 

109 """ 

110 params = conditions_to_filter(query.where) 

111 

112 package_name = params["package"] 

113 mirrors = params.get("mirrors", None) 

114 

115 select_statement_parser = SELECTQueryParser( 

116 query, 

117 PyPIOverallTable.name, 

118 self.get_columns(), 

119 ) 

120 ( 

121 selected_columns, 

122 _, 

123 order_by_conditions, 

124 result_limit, 

125 ) = select_statement_parser.parse_query() 

126 

127 raw_df = self.handler.connection(name=package_name, limit=result_limit).overall( 

128 mirrors=mirrors 

129 ) 

130 

131 select_statement_executor = SELECTQueryExecutor( 

132 raw_df, selected_columns, [], order_by_conditions 

133 ) 

134 

135 result_df = select_statement_executor.execute_query() 

136 

137 return result_df 

138 

139 

140class PyPIPythonMajorTable(CustomAPITable): 

141 name: str = "python_major" 

142 

143 def select(self, query: ast.Select) -> pd.DataFrame: 

144 """triggered at the SELECT query 

145 

146 Args: 

147 query (ast.Select): user's entered query 

148 

149 Returns: 

150 pd.DataFrame: the queried information 

151 """ 

152 params = conditions_to_filter(query.where) 

153 

154 package_name = params["package"] 

155 version = params.get("version", None) 

156 

157 select_statement_parser = SELECTQueryParser( 

158 query, 

159 PyPIOverallTable.name, 

160 self.get_columns(), 

161 ) 

162 ( 

163 selected_columns, 

164 _, 

165 order_by_conditions, 

166 result_limit, 

167 ) = select_statement_parser.parse_query() 

168 

169 raw_df = self.handler.connection( 

170 name=package_name, limit=result_limit 

171 ).python_major(version=version) 

172 

173 select_statement_executor = SELECTQueryExecutor( 

174 raw_df, selected_columns, [], order_by_conditions 

175 ) 

176 

177 result_df = select_statement_executor.execute_query() 

178 

179 return result_df 

180 

181 

182class PyPIPythonMinorTable(CustomAPITable): 

183 name: str = "python_minor" 

184 

185 def select(self, query: ast.Select) -> pd.DataFrame: 

186 """triggered at the SELECT query 

187 

188 Args: 

189 query (ast.Select): user's entered query 

190 

191 Returns: 

192 pd.DataFrame: the queried information 

193 """ 

194 params = conditions_to_filter(query.where) 

195 

196 package_name = params["package"] 

197 version = params.get("version", None) 

198 

199 select_statement_parser = SELECTQueryParser( 

200 query, 

201 PyPIOverallTable.name, 

202 self.get_columns(), 

203 ) 

204 ( 

205 selected_columns, 

206 _, 

207 order_by_conditions, 

208 result_limit, 

209 ) = select_statement_parser.parse_query() 

210 

211 raw_df = self.handler.connection( 

212 name=package_name, limit=result_limit 

213 ).python_minor(version=version) 

214 

215 select_statement_executor = SELECTQueryExecutor( 

216 raw_df, selected_columns, [], order_by_conditions 

217 ) 

218 

219 result_df = select_statement_executor.execute_query() 

220 

221 return result_df 

222 

223 

224class PyPISystemTable(CustomAPITable): 

225 name: str = "system" 

226 

227 def select(self, query: ast.Select) -> pd.DataFrame: 

228 """triggered at the SELECT query 

229 

230 Args: 

231 query (ast.Select): user's entered query 

232 

233 Returns: 

234 pd.DataFrame: the queried information 

235 """ 

236 params = conditions_to_filter(query.where) 

237 

238 package_name = params["package"] 

239 os = params.get("os", None) 

240 

241 select_statement_parser = SELECTQueryParser( 

242 query, 

243 PyPIOverallTable.name, 

244 self.get_columns(), 

245 ) 

246 ( 

247 selected_columns, 

248 _, 

249 order_by_conditions, 

250 result_limit, 

251 ) = select_statement_parser.parse_query() 

252 

253 raw_df = self.handler.connection(name=package_name, limit=result_limit).system( 

254 os=os 

255 ) 

256 

257 select_statement_executor = SELECTQueryExecutor( 

258 raw_df, selected_columns, [], order_by_conditions 

259 ) 

260 

261 result_df = select_statement_executor.execute_query() 

262 

263 return result_df