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 

25from typing import Any, Dict 

26import itertools 

27from numpy import array 

28 

29from .. import Butler 

30from ..core.utils import globToRegex 

31 

32 

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

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

35 

36 Parameters 

37 ---------- 

38 repo : `str` 

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

40 repo and its location. 

41 glob : iterable [`str`] 

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

43 the dataset type names to search for. 

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

45 If provided, only return collections of these types. 

46 chains : `str` 

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

48 Affects contents and formatting of results, see 

49 ``cli.commands.query_collections``. 

50 

51 Returns 

52 ------- 

53 collections : `astropy.table.Table` 

54 A table containing information about collections. 

55 """ 

56 butler = Butler(repo) 

57 expression = globToRegex(glob) 

58 # Only pass expression to queryCollections if there is an expression to 

59 # apply; otherwise let queryCollections use its default value. 

60 kwargs: Dict[str, Any] = {} 

61 if expression: 

62 kwargs["expression"] = expression 

63 

64 if chains == "TABLE": 

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

66 **kwargs)) 

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

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

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

70 

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

72 # populated: 

73 if any(collectionDefinitions): 

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

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

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

77 elif chains == "TREE": 

78 def getCollections(collectionName, nesting=0): 

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

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

81 collection names are indended from their parents by adding spaces 

82 before the collection name. 

83 

84 Parameters 

85 ---------- 

86 collectionName : `str` 

87 The name of the collection to get. 

88 nesting : `int` 

89 The amount of indent to apply before each collection. 

90 

91 Returns 

92 ------- 

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

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

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

96 of its children follows, and so on. 

97 """ 

98 def nested(val): 

99 stepDepth = 2 

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

101 

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

103 if collectionType == "CHAINED": 

104 # Get the child collections of the chained collection: 

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

106 

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

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

109 for child in childCollections]) 

110 

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

112 # the list, and return the list: 

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

114 else: 

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

116 

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

118 **kwargs) 

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

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

121 elif chains == "FLATTEN": 

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

123 flattenChains=True, 

124 **kwargs)) 

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

126 return Table((collectionNames, 

127 collectionTypes), 

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

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