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

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/>.
22from __future__ import annotations
24from astropy.table import Table
25import itertools
26from numpy import array
28from .. import Butler
29from ..core.utils import globToRegex
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 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)]
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.
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.
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
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))
100 # Fill in the child collections of the chained collection:
101 collections = itertools.chain(*[getCollections(child, nesting + 1)
102 for child in childCollections])
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)]
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}")