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 versions: list[IndexItemVersion] = [] 

131 for v in item.versions: 

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

133 

134 def get_bioimageio_yaml_url(base_url: str, bioimageio_yaml: str) -> str: 

135 return base_url.format(bioimageio_yaml=bioimageio_yaml) 

136 

137 url = get_bioimageio_yaml_url( 

138 bioimageio_yaml_base_url, "bioimageio.yaml" 

139 ) 

140 try: 

141 sha256 = _initialize_report_directory(item, v, url) 

142 except Exception: 

143 url = get_bioimageio_yaml_url(bioimageio_yaml_base_url, "rdf.yaml") 

144 sha256 = _initialize_report_directory(item, v, url) 

145 

146 versions.append( 

147 IndexItemVersion( 

148 version=v.version, 

149 comment=v.comment, 

150 created_at=v.created_at, 

151 source=url, 

152 sha256=sha256, 

153 ) 

154 ) 

155 index_items.append(IndexItem(id=item.id, versions=versions, type=item.type)) 

156 

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

158 for item in index_items: 

159 count_per_type[item.type] += 1 

160 

161 index = Index( 

162 items=index_items, 

163 total=len(index_items), 

164 count_per_type=dict(count_per_type), 

165 ) 

166 

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

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

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

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

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

172 

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

174 

175 logger.info( 

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

177 len(index.items), 

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

179 ) 

180 return index 

181 

182 

183def _initialize_report_directory( 

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

185) -> str: 

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

187 

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

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

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

191 _ = r.raise_for_status() 

192 data = r.content 

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

194 

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

196 existing_sha256 = summary.rdf_yaml_sha256 

197 if existing_sha256 == sha256: 

198 logger.info( 

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

200 item.id, 

201 v.version, 

202 sha256, 

203 ) 

204 return sha256 

205 else: 

206 if existing_sha256: 

207 logger.warning( 

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

209 item.id, 

210 v.version, 

211 existing_sha256, 

212 sha256, 

213 ) 

214 if report_path.exists(): 

215 shutil.rmtree(report_path) 

216 

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

218 try: 

219 rdf_content = yaml.load(data) 

220 except Exception as e: 

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

222 

223 summary = InitialSummary( 

224 rdf_content=rdf_content, 

225 rdf_yaml_sha256=sha256, 

226 status="untested", 

227 ) 

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

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

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

231 return sha256 

232 

233 

234if __name__ == "__main__": 

235 _ = create_index()