Coverage for python/lsst/daf/butler/script/collectionChain.py: 10%
52 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-26 02:47 -0700
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-26 02:47 -0700
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30from collections.abc import Iterable
32from .._butler import Butler
33from ..registry import CollectionType, MissingCollectionError
34from ..registry.wildcards import CollectionWildcard
37def collectionChain(
38 repo: str, mode: str, parent: str, children: Iterable[str], doc: str | None, flatten: bool
39) -> tuple[str, ...]:
40 """Get the collections whose names match an expression.
42 Parameters
43 ----------
44 repo : `str`
45 URI to the location of the repo or URI to a config file describing the
46 repo and its location.
47 mode : `str`
48 Update mode for this chain. Options are:
49 'redefine': Create or modify ``parent`` to be defined by the supplied
50 ``children``.
51 'remove': Modify existing chain to remove ``children`` from it.
52 'prepend': Add the given ``children`` to the beginning of the chain.
53 'extend': Modify existing chain to add ``children`` to the end of it.
54 'pop': Pop a numbered element off the chain. Defaults to popping
55 the first element (0). ``children`` must be integers if given.
56 Both 'prepend' and 'extend' are the same as 'redefine' if the chain
57 does not exist.
58 parent : `str`
59 Name of the chained collection to update. Will be created if it
60 does not exist already.
61 children : iterable of `str`
62 Names of the children to be included in the chain.
63 doc : `str`
64 If the chained collection is being created, the documentation string
65 that will be associated with it.
66 flatten : `bool`
67 If `True`, recursively flatten out any nested
68 `~CollectionType.CHAINED` collections in ``children`` first.
70 Returns
71 -------
72 chain : `tuple` of `str`
73 The collections in the chain following this command.
74 """
75 butler = Butler.from_config(repo, writeable=True, without_datastore=True)
77 # Every mode needs children except pop.
78 if not children and mode != "pop":
79 raise RuntimeError(f"Must provide children when defining a collection chain in mode {mode}.")
81 try:
82 butler.registry.getCollectionType(parent)
83 except MissingCollectionError:
84 # Create it -- but only if mode can work with empty chain.
85 if mode in ("redefine", "extend", "prepend"):
86 if not doc:
87 doc = None
88 butler.registry.registerCollection(parent, CollectionType.CHAINED, doc)
89 else:
90 raise RuntimeError(
91 f"Mode '{mode}' requires that the collection exists "
92 f"but collection '{parent}' is not known to this registry"
93 ) from None
95 if flatten:
96 if mode not in ("redefine", "prepend", "extend"):
97 raise RuntimeError(f"'flatten' flag is not allowed for {mode}")
98 wildcard = CollectionWildcard.from_names(children)
99 children = butler.registry.queryCollections(wildcard, flattenChains=True)
101 _modify_collection_chain(butler, mode, parent, children)
103 return tuple(butler.registry.getCollectionChain(parent))
106def _modify_collection_chain(butler: Butler, mode: str, parent: str, children: Iterable[str]) -> None:
107 if mode == "prepend":
108 butler.collection_chains.prepend_chain(parent, children)
109 elif mode == "redefine":
110 butler.collection_chains.redefine_chain(parent, children)
111 elif mode == "remove":
112 butler.collection_chains.remove_from_chain(parent, children)
113 elif mode == "pop":
114 children_to_pop = _find_children_to_pop(butler, parent, children)
115 butler.collection_chains.remove_from_chain(parent, children_to_pop)
116 elif mode == "extend":
117 butler.collection_chains.extend_chain(parent, children)
118 else:
119 raise ValueError(f"Unrecognized update mode: '{mode}'")
122def _find_children_to_pop(butler: Butler, parent: str, children: Iterable[str]) -> list[str]:
123 """Find the names of the children in the parent collection corresponding to
124 the given indexes.
125 """
126 children = list(children)
127 current = butler.registry.getCollectionChain(parent)
128 n_current = len(current)
129 if children:
131 def convert_index(i: int) -> int:
132 """Convert negative index to positive."""
133 if i >= 0:
134 return i
135 return n_current + i
137 # For this mode the children should be integers.
138 # Convert negative integers to positive ones to allow
139 # sorting.
140 indices = [convert_index(int(child)) for child in children]
141 else:
142 # Nothing specified, pop from the front of the chain.
143 indices = [0]
145 for index in indices:
146 if index >= n_current:
147 raise IndexError(
148 f"Index {index} is out of range. Parent collection {parent} has {n_current} children."
149 )
151 return [current[i] for i in indices]