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
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
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/>.
22from __future__ import annotations
24import itertools
26from astropy.table import Table
27from numpy import array
29from .._butler import Butler
32def queryCollections(repo, glob, collection_type, chains):
33 """Get the collections whose names match an expression.
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``.
50 Returns
51 -------
52 collections : `astropy.table.Table`
53 A table containing information about collections.
54 """
55 butler = Butler(repo)
57 if not glob:
58 glob = ...
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 ]
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":
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.
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.
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 """
101 def nested(val):
102 stepDepth = 2
103 return " " * (stepDepth * nesting) + val
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))
110 # Fill in the child collections of the chained collection:
111 collections = itertools.chain(
112 *[getCollections(child, nesting + 1) for child in childCollections]
113 )
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)]
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}")