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
25from typing import Any, Dict
26import itertools
27from numpy import array
29from .. import Butler
30from ..core.utils import globToRegex
33def queryCollections(repo, glob, collection_type, chains):
34 """Get the collections whose names match an expression.
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``.
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
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)]
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.
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.
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
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))
107 # Fill in the child collections of the chained collection:
108 collections = itertools.chain(*[getCollections(child, nesting + 1)
109 for child in childCollections])
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)]
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}")