Coverage for python/lsst/daf/butler/script/queryCollections.py: 11%
35 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-01 19:55 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-01 19:55 +0000
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
31def queryCollections(repo, glob, collection_type, chains):
32 """Get the collections whose names match an expression.
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``.
49 Returns
50 -------
51 collections : `astropy.table.Table`
52 A table containing information about collections.
53 """
54 butler = Butler(repo)
56 if not glob:
57 glob = ...
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)]
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.
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.
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
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))
102 # Fill in the child collections of the chained collection:
103 collections = itertools.chain(*[getCollections(child, nesting + 1)
104 for child in childCollections])
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)]
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}")