Coverage for src/backoffice/index.py: 0%

1"""Data models and functions for indexing the bioimage.io collection""" 

2 

3from __future__ import annotations 

4 

5import hashlib 

6import json 

7import shutil 

8from collections import defaultdict 

9from collections.abc import Sequence 

10from datetime import datetime 

11from pathlib import Path 

12 

13try: 

14 import httpx 

15 from loguru import logger 

16 from pydantic import BaseModel, Field 

17 

18 from backoffice._settings import settings 

19 from backoffice.compatibility import InitialSummary 

20 from backoffice.utils import ( 

21 get_summary, 

22 get_summary_file_path, 

23 yaml, 

24 ) 

25except ImportError as e: 

26 raise ImportError( 

27 "Missing dependencies. " 

28 "Please install `backoffice[dev]` to use backoffice.index." 

29 ) from e 

30 

31 

32from .utils_pure import get_report_path 

33 

34 

35class Node(BaseModel, frozen=True, extra="ignore"): 

36 pass 

37 

38 

39class ResponseItemVersion(Node, frozen=True): 

40 version: str 

41 comment: str | None 

42 created_at: datetime 

43 

44 

45class IndexItemVersion(Node, frozen=True): 

46 version: str 

47 comment: str | None 

48 created_at: datetime 

49 source: str 

50 sha256: str 

51 

52 

53class ResponseItem(Node, frozen=True): 

54 id: str 

55 versions: Sequence[ResponseItemVersion] 

56 type: str 

57 

58 

59class IndexItem(Node, frozen=True): 

60 id: str 

61 versions: Sequence[IndexItemVersion] 

62 type: str 

63 

64 

65class Response(Node, frozen=True): 

66 """Response from Hypha list endpoint""" 

67 

68 items: list[ResponseItem] 

69 total: int 

70 offset: int 

71 limit: int 

72 

73 

74class Index(Node, frozen=True): 

75 items: list[IndexItem] 

76 total: int 

77 count_per_type: dict[str, int] 

78 timestamp: datetime = Field(default_factory=datetime.now) 

79 

80 

81def load_index(path: Path = Path("index.json")) -> Index: 

82 logger.info("loading index from {}", path) 

83 return Index.model_validate_json(path.read_text(encoding="utf-8")) 

84 

85 

86def create_index() -> Index: 

87 """Index the bioimage.io collection""" 

88 

89 index_path = Path("index.json") 

90 if index_path.exists(): 

91 index = load_index(index_path) 

92 else: 

93 url = f"{settings.hypha_base_url}/public/services/artifact-manager/list" 

94 

95 def request(offset: int) -> Response: 

96 r = httpx.get( 

97 url, 

98 params={ 

99 "parent_id": "bioimage-io/bioimage.io", 

100 "offset": offset, 

101 "pagination": True, 

102 "limit": 10000, 

103 }, 

104 headers=settings.get_hypha_headers(), 

105 timeout=settings.http_timeout, 

106 ) 

107 try: 

108 _ = r.raise_for_status() 

109 except Exception: 

110 logger.error(r.json()) 

111 raise 

112 else: 

113 return Response.model_validate_json(r.content) 

114 

115 items: list[ResponseItem] = [] 

116 for page in range(100): 

117 response = request(len(items)) 

118 logger.info("Page {}: {} entries", page, len(response.items)) 

119 items.extend(response.items) 

120 if response.total <= len(items): 

121 if response.total != len(items): 

122 logger.error( 

123 "response.total {} != len(items) {}", response.total, len(items) 

124 ) 

125 break 

126 

127 index_items: list[IndexItem] = [] 

128 for item in items: 

129 domain, item_id_wo_domain = item.id.split("/", 1) 

130 try: 

131 item_id = item.id 

132 item_type = item.type 

133 versions: list[IndexItemVersion] = [] 

134 for v in item.versions: 

135 bioimageio_yaml_base_url = f"{settings.hypha_base_url}/{domain}/artifacts/{item_id_wo_domain}/files/{{bioimageio_yaml}}?version={v.version}" 

136 

137 def get_bioimageio_yaml_url( 

138 base_url: str, bioimageio_yaml: str 

139 ) -> str: 

140 return base_url.format(bioimageio_yaml=bioimageio_yaml) 

141 

142 url = get_bioimageio_yaml_url( 

143 bioimageio_yaml_base_url, "bioimageio.yaml" 

144 ) 

145 try: 

146 sha256 = _initialize_report_directory(item, v, url) 

147 except Exception: 

148 url = get_bioimageio_yaml_url( 

149 bioimageio_yaml_base_url, "rdf.yaml" 

150 ) 

151 sha256 = _initialize_report_directory(item, v, url) 

152 

153 versions.append( 

154 IndexItemVersion( 

155 version=v.version, 

156 comment=v.comment, 

157 created_at=v.created_at, 

158 source=url, 

159 sha256=sha256, 

160 ) 

161 ) 

162 index_items.append( 

163 IndexItem(id=item_id, versions=versions, type=item_type) 

164 ) 

165 except Exception as e: 

166 logger.error( 

167 "Error processing index item {}: {}. Skipping this item.", 

168 item, 

169 str(e), 

170 ) 

171 

172 count_per_type = defaultdict[str, int](int) 

173 for item in index_items: 

174 count_per_type[item.type] += 1 

175 

176 index = Index( 

177 items=index_items, 

178 total=len(index_items), 

179 count_per_type=dict(count_per_type), 

180 ) 

181 

182 json_dict = index.model_dump(mode="json") 

183 with index_path.open("wt", encoding="utf-8") as f: 

184 json.dump(json_dict, f, indent=4, sort_keys=True, ensure_ascii=False) 

185 # TODO: use .model_dump_json once it supports 'sort_keys' argument for a potential speed gain 

186 # _ = index_path.write_text(index.model_dump_json(indent=4), encoding="utf-8") 

187 

188 logger.info("saved index to {}", index_path) 

189 

190 logger.info( 

191 "loaded index with {} ids and {} versions", 

192 len(index.items), 

193 sum(len(item.versions) for item in index.items), 

194 ) 

195 return index 

196 

197 

198def _initialize_report_directory( 

199 item: ResponseItem, v: ResponseItemVersion, url: str 

200) -> str: 

201 """Initialize the report directory for an item version. 

202 

203 Returns sha256 of the rdf.yaml file.""" 

204 report_path = get_report_path(item.id, v.version) 

205 r = httpx.get(url, follow_redirects=True, timeout=settings.http_timeout) 

206 _ = r.raise_for_status() 

207 data = r.content 

208 sha256 = hashlib.sha256(data).hexdigest() 

209 

210 summary = get_summary(item.id, v.version) 

211 existing_sha256 = summary.rdf_yaml_sha256 

212 if existing_sha256 == sha256: 

213 logger.info( 

214 "Found existing summary for {}/{} with matching RDF SHA-256: {}", 

215 item.id, 

216 v.version, 

217 sha256, 

218 ) 

219 return sha256 

220 else: 

221 if existing_sha256: 

222 logger.warning( 

223 "Found existing summary for {}/{} with different RDF SHA-256: {} != {}. deleting and replacing...", 

224 item.id, 

225 v.version, 

226 existing_sha256, 

227 sha256, 

228 ) 

229 if report_path.exists(): 

230 shutil.rmtree(report_path) 

231 

232 report_path.mkdir(parents=True, exist_ok=True) 

233 try: 

234 rdf_content = yaml.load(data) 

235 except Exception as e: 

236 rdf_content = {"error": str(e)} 

237 

238 summary = InitialSummary( 

239 rdf_content=rdf_content, 

240 rdf_yaml_sha256=sha256, 

241 status="untested", 

242 ) 

243 summary_path = get_summary_file_path(item.id, v.version) 

244 _ = summary_path.write_text(summary.model_dump_json(indent=4), encoding="utf-8") 

245 logger.info("Initialized report directory {}", report_path) 

246 return sha256 

247 

248 

249if __name__ == "__main__": 

250 _ = create_index()