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 

24from astropy.table import Table 

25import itertools 

26from numpy import array 

27 

28from .. import Butler 

29 

30 

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

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

33 

34 Parameters 

35 ---------- 

36 repo : `str` 

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

38 repo and its location. 

39 glob : iterable [`str`] 

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

41 the dataset type names to search for. 

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

43 If provided, only return collections of these types. 

44 chains : `str` 

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

46 Affects contents and formatting of results, see 

47 ``cli.commands.query_collections``. 

48 

49 Returns 

50 ------- 

51 collections : `astropy.table.Table` 

52 A table containing information about collections. 

53 """ 

54 butler = Butler(repo) 

55 

56 if not glob: 

57 glob = ... 

58 

59 if chains == "TABLE": 

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

61 expression=glob)) 

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

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

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

65 

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

67 # populated: 

68 if any(collectionDefinitions): 

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

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

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

72 elif chains == "TREE": 

73 def getCollections(collectionName, nesting=0): 

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

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

76 collection names are indended from their parents by adding spaces 

77 before the collection name. 

78 

79 Parameters 

80 ---------- 

81 collectionName : `str` 

82 The name of the collection to get. 

83 nesting : `int` 

84 The amount of indent to apply before each collection. 

85 

86 Returns 

87 ------- 

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

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

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

91 of its children follows, and so on. 

92 """ 

93 def nested(val): 

94 stepDepth = 2 

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

96 

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

98 if collectionType == "CHAINED": 

99 # Get the child collections of the chained collection: 

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

101 

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

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

104 for child in childCollections]) 

105 

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

107 # the list, and return the list: 

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

109 else: 

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

111 

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

113 expression=glob) 

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

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

116 elif chains == "FLATTEN": 

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

118 flattenChains=True, 

119 expression=glob)) 

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

121 return Table((collectionNames, 

122 collectionTypes), 

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

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