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

Shortcuts 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

35 statements  

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 

24import itertools 

25 

26from astropy.table import Table 

27from numpy import array 

28 

29from .._butler import Butler 

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 not glob: 

58 glob = ... 

59 

60 if chains == "TABLE": 

61 collectionNames = list( 

62 butler.registry.queryCollections(collectionTypes=frozenset(collection_type), expression=glob) 

63 ) 

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

65 collectionDefinitions = [ 

66 str(butler.registry.getCollectionChain(name)) if colType == "CHAINED" else "" 

67 for name, colType in zip(collectionNames, collectionTypes) 

68 ] 

69 

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

71 # populated: 

72 if any(collectionDefinitions): 

73 return Table( 

74 (collectionNames, collectionTypes, collectionDefinitions), 

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

76 ) 

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

78 elif chains == "TREE": 

79 

80 def getCollections(collectionName, nesting=0): 

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

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

83 collection names are indended from their parents by adding spaces 

84 before the collection name. 

85 

86 Parameters 

87 ---------- 

88 collectionName : `str` 

89 The name of the collection to get. 

90 nesting : `int` 

91 The amount of indent to apply before each collection. 

92 

93 Returns 

94 ------- 

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

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

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

98 of its children follows, and so on. 

99 """ 

100 

101 def nested(val): 

102 stepDepth = 2 

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

104 

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

106 if collectionType == "CHAINED": 

107 # Get the child collections of the chained collection: 

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

109 

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

111 collections = itertools.chain( 

112 *[getCollections(child, nesting + 1) for child in childCollections] 

113 ) 

114 

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

116 # the list, and return the list: 

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

118 else: 

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

120 

121 collectionNameIter = butler.registry.queryCollections( 

122 collectionTypes=frozenset(collection_type), expression=glob 

123 ) 

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

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

126 elif chains == "FLATTEN": 

127 collectionNames = list( 

128 butler.registry.queryCollections( 

129 collectionTypes=frozenset(collection_type), flattenChains=True, expression=glob 

130 ) 

131 ) 

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

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

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