Coverage for python/lsst/daf/butler/script/queryCollections.py: 7%

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

28from __future__ import annotations 

29 

30from collections.abc import Iterable 

31from fnmatch import fnmatch 

32from typing import Literal 

33 

34from astropy.table import Column, Table, hstack, vstack 

35 

36from .._butler import Butler 

37from .._butler_collections import CollectionInfo 

38from .._collection_type import CollectionType 

39 

40 

41def _parseDatasetTypes(dataset_types: frozenset[str] | list[str] | None) -> list[str]: 

42 """Parse dataset types from a collection info object or a list of strings. 

43 

44 Parameters 

45 ---------- 

46 dataset_types : `frozenset`[`str`] | `list`[`str`] | `None` 

47 The dataset types to parse. If `None`, an empty list is returned. 

48 If a `frozenset` or `list` is provided, it is returned as a list. 

49 

50 Returns 

51 ------- 

52 dataset_types : `list`[`str`] 

53 The parsed dataset types. 

54 """ 

55 return [""] if not dataset_types else list(dataset_types) 

56 

57 

58def _getTable( 

59 repo: str, 

60 glob: Iterable[str], 

61 collection_type: Iterable[CollectionType], 

62 inverse: bool, 

63 show_dataset_types: bool = False, 

64 exclude_dataset_types: Iterable[str] | None = None, 

65) -> Table: 

66 """Run queryCollections and return the results in Table form. 

67 

68 Only lists the first child (or parent if `inverse` is `True`) in the 

69 description column. 

70 

71 Parameters 

72 ---------- 

73 repo : `str` 

74 The Butler repository location. 

75 glob : `collections.abc.Iterable` of `str` 

76 Wildcard to pass to ``queryCollections``. 

77 collection_type 

78 Same as `queryCollections` 

79 inverse : `bool` 

80 True if parent CHAINED datasets of each dataset should be listed in the 

81 description column, False if children of CHAINED datasets should be 

82 listed. 

83 show_dataset_types : `bool`, optional 

84 If `True`, also show the dataset types present within each collection. 

85 exclude_dataset_types : `~collections.abc.Iterable` [ `str` ], optional 

86 A glob-style iterable of dataset types to exclude. 

87 Only has an effect if `show_dataset_types` is True. 

88 

89 Returns 

90 ------- 

91 collections : `astropy.table.Table` 

92 Same as `queryCollections` 

93 """ 

94 typeCol = "Type" 

95 descriptionCol = "Parents" if inverse else "Children" 

96 table = Table( 

97 names=("Name", typeCol, descriptionCol), 

98 dtype=(str, str, str), 

99 ) 

100 if show_dataset_types: 

101 table.add_column(Column(name="Dataset Types", dtype=str)) 

102 butler = Butler.from_config(repo) 

103 

104 def addDatasetTypes(collection_table: Table, collection: str, dataset_types: list[str]) -> Table: 

105 if dataset_types[0] == "": 

106 cinfo = butler.collections.get_info(collection, include_summary=True) 

107 dataset_types = _parseDatasetTypes(cinfo.dataset_types) 

108 if exclude_dataset_types: 

109 dataset_types = [ 

110 dt 

111 for dt in dataset_types 

112 if not any(fnmatch(dt, pattern) for pattern in exclude_dataset_types) 

113 ] 

114 dataset_types = _parseDatasetTypes(dataset_types) 

115 types_table = Table({"Dataset Types": sorted(dataset_types)}, dtype=(str,)) 

116 collection_table = hstack([collection_table, types_table]).filled("") 

117 return collection_table 

118 

119 def addCollection(info: CollectionInfo, relation: str) -> None: 

120 try: 

121 info_relatives = getattr(info, relation) 

122 except AttributeError: 

123 info_relatives = [] 

124 # Parent results can be returned in a non-deterministic order, so sort 

125 # them to make the output deterministic. 

126 if relation == "parents": 

127 info_relatives = sorted(info_relatives) 

128 if info_relatives: 

129 collection_table = Table([[info.name], [info.type.name]], names=("Name", typeCol)) 

130 description_table = Table(names=(descriptionCol,), dtype=(str,)) 

131 for info_relative in info_relatives: 

132 relative_table = Table([[info_relative]], names=(descriptionCol,)) 

133 if show_dataset_types: 

134 relative_table = addDatasetTypes(relative_table, info_relative, [""]) 

135 description_table = vstack([description_table, relative_table]) 

136 collection_table = hstack([collection_table, description_table]).filled("") 

137 for row in collection_table: 

138 table.add_row(row) 

139 else: 

140 collection_table = Table( 

141 [[info.name], [info.type.name], [""]], names=("Name", typeCol, descriptionCol) 

142 ) 

143 if show_dataset_types: 

144 collection_table = addDatasetTypes(collection_table, info.name, [""]) 

145 for row in collection_table: 

146 table.add_row(row) 

147 

148 collections = sorted( 

149 butler.collections.query_info( 

150 glob or "*", 

151 collection_types=frozenset(collection_type), 

152 include_parents=inverse, 

153 include_summary=show_dataset_types, 

154 ) 

155 ) 

156 if inverse: 

157 for info in collections: 

158 addCollection(info, "parents") 

159 # If none of the datasets has a parent dataset then remove the 

160 # description column. 

161 if not any(c for c in table[descriptionCol]): 

162 del table[descriptionCol] 

163 else: 

164 for info in collections: 

165 if info.type == CollectionType.CHAINED: 

166 addCollection(info, "children") 

167 else: 

168 addCollection(info, "self") 

169 # If there aren't any CHAINED datasets in the results then remove the 

170 # description column. 

171 if not any(columnVal == CollectionType.CHAINED.name for columnVal in table[typeCol]): 

172 del table[descriptionCol] 

173 

174 return table 

175 

176 

177def _getTree( 

178 repo: str, 

179 glob: Iterable[str], 

180 collection_type: Iterable[CollectionType], 

181 inverse: bool, 

182 show_dataset_types: bool = False, 

183 exclude_dataset_types: Iterable[str] | None = None, 

184) -> Table: 

185 """Run queryCollections and return the results in a table representing tree 

186 form. 

187 

188 Recursively lists children (or parents if `inverse` is `True`) 

189 

190 Parameters 

191 ---------- 

192 repo : `str` 

193 Butler repository location. 

194 glob : `collections.abc.Iterable` of `str` 

195 Wildcards to pass to ``queryCollections``. 

196 collection_type 

197 Same as `queryCollections` 

198 inverse : `bool` 

199 True if parent CHAINED datasets of each dataset should be listed in the 

200 description column, False if children of CHAINED datasets should be 

201 listed. 

202 show_dataset_types : `bool`, optional 

203 If `True`, also show the dataset types present within each collection. 

204 exclude_dataset_types : `~collections.abc.Iterable` [ `str` ], optional 

205 A glob-style iterable of dataset types to exclude. 

206 Only has an effect if `show_dataset_types` is True. 

207 

208 Returns 

209 ------- 

210 collections : `astropy.table.Table` 

211 Same as `queryCollections` 

212 """ 

213 table = Table( 

214 names=("Name", "Type"), 

215 dtype=(str, str), 

216 ) 

217 if show_dataset_types: 

218 table.add_column(Column(name="Dataset Types", dtype=str)) 

219 butler = Butler.from_config(repo, without_datastore=True) 

220 

221 def addCollection(info: CollectionInfo, level: int = 0) -> None: 

222 collection_table = Table([[" " * level + info.name], [info.type.name]], names=["Name", "Type"]) 

223 if show_dataset_types: 

224 if info.type == CollectionType.CHAINED: 

225 collection_table = hstack( 

226 [collection_table, Table([[""] * len(collection_table)], names=["Dataset Types"])] 

227 ) 

228 else: 

229 dataset_types = _parseDatasetTypes(info.dataset_types) 

230 if exclude_dataset_types: 

231 dataset_types = [ 

232 dt 

233 for dt in dataset_types 

234 if not any(fnmatch(dt, pattern) for pattern in exclude_dataset_types) 

235 ] 

236 dataset_types = _parseDatasetTypes(dataset_types) 

237 dataset_types_table = Table({"Dataset Types": sorted(dataset_types)}, dtype=(str,)) 

238 collection_table = hstack([collection_table, dataset_types_table]).filled("") 

239 for row in collection_table: 

240 table.add_row(row) 

241 

242 if inverse: 

243 assert info.parents is not None # For mypy. 

244 for pname in sorted(info.parents): 

245 pinfo = butler.collections.get_info( 

246 pname, include_parents=inverse, include_summary=show_dataset_types 

247 ) 

248 addCollection(pinfo, level + 1) 

249 else: 

250 if info.type == CollectionType.CHAINED: 

251 for name in info.children: 

252 cinfo = butler.collections.get_info(name, include_summary=show_dataset_types) 

253 addCollection(cinfo, level + 1) 

254 

255 collections = butler.collections.query_info( 

256 glob or "*", 

257 collection_types=frozenset(collection_type), 

258 include_parents=inverse, 

259 include_summary=show_dataset_types, 

260 ) 

261 for collection in sorted(collections): 

262 addCollection(collection) 

263 return table 

264 

265 

266def _getList( 

267 repo: str, 

268 glob: Iterable[str], 

269 collection_type: Iterable[CollectionType], 

270 flatten_chains: bool, 

271 show_dataset_types: bool = False, 

272 exclude_dataset_types: Iterable[str] | None = None, 

273) -> Table: 

274 """Return collection results as a table representing a flat list of 

275 collections. 

276 

277 Parameters 

278 ---------- 

279 repo : `str` 

280 Butler repository location. 

281 glob : `collections.abc.Iterable` of `str` 

282 Wildcards to pass to ``queryCollections``. 

283 collection_type 

284 Same as `queryCollections` 

285 flatten_chains : `bool` 

286 If `True`, flatten the tree of CHAINED datasets. 

287 show_dataset_types : `bool`, optional 

288 If `True`, also show the dataset types present within each collection. 

289 exclude_dataset_types : `~collections.abc.Iterable` [ `str` ], optional 

290 A glob-style iterable of dataset types to exclude. 

291 Only has an effect if `show_dataset_types` is True. 

292 

293 Returns 

294 ------- 

295 collections : `astropy.table.Table` 

296 Same as `queryCollections` 

297 """ 

298 table = Table( 

299 names=("Name", "Type"), 

300 dtype=(str, str), 

301 ) 

302 if show_dataset_types: 

303 table.add_column(Column(name="Dataset Types", dtype=str)) 

304 butler = Butler.from_config(repo) 

305 

306 def addCollection(info: CollectionInfo) -> None: 

307 collection_table = Table([[info.name], [info.type.name]], names=["Name", "Type"]) 

308 if show_dataset_types: 

309 dataset_types = _parseDatasetTypes(info.dataset_types) 

310 if exclude_dataset_types: 

311 dataset_types = [ 

312 dt 

313 for dt in dataset_types 

314 if not any(fnmatch(dt, pattern) for pattern in exclude_dataset_types) 

315 ] 

316 dataset_types = _parseDatasetTypes(dataset_types) 

317 dataset_types_table = Table({"Dataset Types": sorted(dataset_types)}, dtype=(str,)) 

318 collection_table = hstack([collection_table, dataset_types_table]).filled("") 

319 for row in collection_table: 

320 table.add_row(row) 

321 

322 collections = list( 

323 butler.collections.query_info( 

324 glob or "*", 

325 collection_types=frozenset(collection_type), 

326 flatten_chains=flatten_chains, 

327 include_summary=show_dataset_types, 

328 ) 

329 ) 

330 for collection in collections: 

331 addCollection(collection) 

332 

333 return table 

334 

335 

336def queryCollections( 

337 repo: str, 

338 glob: Iterable[str], 

339 collection_type: Iterable[CollectionType], 

340 chains: Literal["INVERSE-TABLE", "TABLE", "TREE", "INVERSE-TREE", "FLATTEN", "NO-CHILDREN"], 

341 show_dataset_types: bool = False, 

342 exclude_dataset_types: Iterable[str] | None = None, 

343) -> Table: 

344 """Get the collections whose names match an expression. 

345 

346 Parameters 

347 ---------- 

348 repo : `str` 

349 URI to the location of the repo or URI to a config file describing the 

350 repo and its location. 

351 glob : `~collections.abc.Iterable` [`str`] 

352 A list of glob-style search string that fully or partially identify 

353 the dataset type names to search for. 

354 collection_type : `~collections.abc.Iterable` [ `CollectionType` ], \ 

355 optional 

356 If provided, only return collections of these types. 

357 chains : `str` 

358 Affects contents and formatting of results, see 

359 ``cli.commands.query_collections``. 

360 show_dataset_types : `bool`, optional 

361 If `True`, include the dataset types present within each collection. 

362 exclude_dataset_types : `~collections.abc.Iterable` [ `str` ], optional 

363 A glob-style iterable of dataset types to exclude. 

364 Only has an effect if `show_dataset_types` is True. 

365 

366 Returns 

367 ------- 

368 collections : `astropy.table.Table` 

369 A table containing information about collections. 

370 """ 

371 if (inverse := chains == "INVERSE-TABLE") or chains == "TABLE": 

372 return _getTable(repo, glob, collection_type, inverse, show_dataset_types, exclude_dataset_types) 

373 elif (inverse := chains == "INVERSE-TREE") or chains == "TREE": 

374 return _getTree(repo, glob, collection_type, inverse, show_dataset_types, exclude_dataset_types) 

375 elif chains == "FLATTEN" or chains == "NO-CHILDREN": 

376 flatten = chains == "FLATTEN" 

377 return _getList(repo, glob, collection_type, flatten, show_dataset_types, exclude_dataset_types) 

378 raise RuntimeError(f"Value for --chains not recognized: {chains}")