Coverage for python/lsst/daf/butler/script/collectionChain.py: 9%
52 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 10:00 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 10:00 +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 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.prepend_collection_chain(parent, children)
109 elif mode == "redefine":
110 butler.registry.setCollectionChain(parent, children)
111 else:
112 current = list(butler.registry.getCollectionChain(parent))
114 if mode == "extend":
115 current.extend(children)
116 children = current
117 elif mode == "remove":
118 for child in children:
119 current.remove(child)
120 children = current
121 elif mode == "pop":
122 if children:
123 n_current = len(current)
125 def convert_index(i: int) -> int:
126 """Convert negative index to positive."""
127 if i >= 0:
128 return i
129 return n_current + i
131 # For this mode the children should be integers.
132 # Convert negative integers to positive ones to allow
133 # sorting.
134 indices = [convert_index(int(child)) for child in children]
136 # Reverse sort order so we can remove from the end first
137 indices = sorted(indices, reverse=True)
139 else:
140 # Nothing specified, pop from the front of the chain.
141 indices = [0]
143 for i in indices:
144 current.pop(i)
146 children = current
147 else:
148 raise ValueError(f"Unrecognized update mode: '{mode}'")
150 butler.registry.setCollectionChain(parent, children)