Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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 program is free software: you can redistribute it and/or modify 

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

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

12# (at your option) any later version. 

13# 

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

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

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

17# GNU General Public License for more details. 

18# 

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

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

21 

22from __future__ import annotations 

23 

24from astropy.table import Table 

25import itertools 

26from numpy import array 

27 

28from .. import Butler 

29from ..core.utils import globToRegex 

30 

31 

32def queryCollections(repo, glob, collection_type, chains): 

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

34 

35 Parameters 

36 ---------- 

37 repo : `str` 

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

39 repo and its location. 

40 glob : iterable [`str`] 

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

42 the dataset type names to search for. 

43 collection_type : `Iterable` [ `CollectionType` ], optional 

44 If provided, only return collections of these types. 

45 chains : `str` 

46 Must be one of "FLATTEN", "TABLE", or "TREE" (case sensitive). 

47 Affects contents and formatting of results, see 

48 ``cli.commands.query_collections``. 

49 

50 Returns 

51 ------- 

52 collections : `astropy.table.Table` 

53 A table containing information about collections. 

54 """ 

55 butler = Butler(repo) 

56 

57 if chains == "TABLE": 

58 collectionNames = list(butler.registry.queryCollections(collectionTypes=frozenset(collection_type), 

59 expression=globToRegex(glob))) 

60 collectionTypes = [butler.registry.getCollectionType(c).name for c in collectionNames] 

61 collectionDefinitions = [str(butler.registry.getCollectionChain(name)) if colType == "CHAINED" else "" 

62 for name, colType in zip(collectionNames, collectionTypes)] 

63 

64 # Only add a definition column if at least one definition row is 

65 # populated: 

66 if any(collectionDefinitions): 

67 return Table((collectionNames, collectionTypes, collectionDefinitions), 

68 names=("Name", "Type", "Definition")) 

69 return Table((collectionNames, collectionTypes), names=("Name", "Type")) 

70 elif chains == "TREE": 

71 def getCollections(collectionName, nesting=0): 

72 """Get a list of the name and type of the passed-in collection, 

73 and its child collections, if it is a CHAINED collection. Child 

74 collection names are indended from their parents by adding spaces 

75 before the collection name. 

76 

77 Parameters 

78 ---------- 

79 collectionName : `str` 

80 The name of the collection to get. 

81 nesting : `int` 

82 The amount of indent to apply before each collection. 

83 

84 Returns 

85 ------- 

86 collections : `list` [`tuple` [`str`, `str`]] 

87 Tuples of the collection name and its type. Starts with the 

88 passed-in collection, and if it is a CHAINED collection, each 

89 of its children follows, and so on. 

90 """ 

91 def nested(val): 

92 stepDepth = 2 

93 return " " * (stepDepth * nesting) + val 

94 

95 collectionType = butler.registry.getCollectionType(collectionName).name 

96 if collectionType == "CHAINED": 

97 # Get the child collections of the chained collection: 

98 childCollections = list(butler.registry.getCollectionChain(collectionName)) 

99 

100 # Fill in the child collections of the chained collection: 

101 collections = itertools.chain(*[getCollections(child, nesting + 1) 

102 for child in childCollections]) 

103 

104 # Insert the chained (parent) collection at the beginning of 

105 # the list, and return the list: 

106 return [(nested(collectionName), "CHAINED")] + list(collections) 

107 else: 

108 return [(nested(collectionName), collectionType)] 

109 

110 collectionNameIter = butler.registry.queryCollections(collectionTypes=frozenset(collection_type), 

111 expression=globToRegex(glob)) 

112 collections = list(itertools.chain(*[getCollections(name) for name in collectionNameIter])) 

113 return Table(array(collections), names=("Name", "Type")) 

114 elif chains == "FLATTEN": 

115 collectionNames = list(butler.registry.queryCollections(collectionTypes=frozenset(collection_type), 

116 flattenChains=True, 

117 expression=globToRegex(glob))) 

118 collectionTypes = [butler.registry.getCollectionType(c).name for c in collectionNames] 

119 return Table((collectionNames, 

120 collectionTypes), 

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

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