Merge pull request #533 from lsst/tickets/DM-30373

DM-30373: Add butler collection-chain command line

doc/changes/DM-30373.feature.rst

@@ -0,0 +1 @@
+Add new ``butler collection-chain`` subcommand for creating collection chains from the command line.

python/lsst/daf/butler/cli/butler.py

@@ -22,6 +22,7 @@
 import abc
 import click
 from collections import defaultdict
+import functools
 import logging
 import os
 import traceback
@@ -36,6 +37,35 @@
 log = logging.getLogger(__name__)
 
 
+@functools.lru_cache
+def _importPlugin(pluginName):
+    """Import a plugin that contains Click commands.
+
+    Parameters
+    ----------
+    pluginName : `str`
+        An importable module whose __all__ parameter contains the commands
+        that can be called.
+
+    Returns
+    -------
+    An imported module or None
+        The imported module, or None if the module could not be imported.
+
+    Notes
+    -----
+    A cache is used in order to prevent repeated reports of failure
+    to import a module that can be triggered by ``butler --help``.
+    """
+    try:
+        return doImport(pluginName)
+    except Exception as err:
+        log.warning("Could not import plugin from %s, skipping.", pluginName)
+        log.debug("Plugin import exception: %s\nTraceback:\n%s", err,
+                  "".join(traceback.format_tb(err.__traceback__)))
+        return None
+
+
 class LoaderCLI(click.MultiCommand, abc.ABC):
     """Extends `click.MultiCommand`, which dispatches to subcommands, to load
     subcommands at runtime."""
@@ -69,10 +99,10 @@ def getLocalCommands(self):
             The key is the command name. The value is a list of package(s) that
             contains the command.
         """
-        commandsLocation = self._importPlugin(self.localCmdPkg)
+        commandsLocation = _importPlugin(self.localCmdPkg)
         if commandsLocation is None:
             # _importPlugins logs an error, don't need to do it again here.
-            return {}
+            return defaultdict(list)
         return defaultdict(list, {self._funcNameToCmdName(f):
                                   [self.localCmdPkg] for f in commandsLocation.__all__})
 
@@ -115,7 +145,7 @@ def get_command(self, ctx, name):
         if name not in commands:
             return None
         self._raiseIfDuplicateCommands(commands)
-        return self._importPlugin(commands[name][0] + "." + self._cmdNameToFuncName(name))
+        return _importPlugin(commands[name][0] + "." + self._cmdNameToFuncName(name))
 
     def _setupLogging(self, ctx):
         """Init the logging system and config it for the command.
@@ -167,29 +197,6 @@ def _cmdNameToFuncName(cls, commandName):
         legal, function name to the command name."""
         return commandName.replace("-", "_")
 
-    @staticmethod
-    def _importPlugin(pluginName):
-        """Import a plugin that contains Click commands.
-
-        Parameters
-        ----------
-        pluginName : `str`
-            An importable module whose __all__ parameter contains the commands
-            that can be called.
-
-        Returns
-        -------
-        An imported module or None
-            The imported module, or None if the module could not be imported.
-        """
-        try:
-            return doImport(pluginName)
-        except Exception as err:
-            log.warning("Could not import plugin from %s, skipping.", pluginName)
-            log.debug("Plugin import exception: %s\nTraceback:\n%s", err,
-                      "".join(traceback.format_tb(err.__traceback__)))
-            return None
-
     @staticmethod
     def _mergeCommandLists(a, b):
         """Combine two dicts whose keys are strings (command name) and values

python/lsst/daf/butler/cli/cmd/__init__.py

@@ -23,6 +23,7 @@
            "butler_import",
            "certify_calibrations",
            "create",
+           "collection_chain",
            "config_dump",
            "config_validate",
            "prune_collection",
@@ -42,6 +43,7 @@
                        butler_import,
                        certify_calibrations,
                        create,
+                       collection_chain,
                        config_dump,
                        config_validate,
                        prune_collection,

python/lsst/daf/butler/cli/cmd/commands.py

@@ -497,3 +497,45 @@ def transfer_datasets(**kwargs):
     """
     number = script.transferDatasets(**kwargs)
     print(f"Number of datasets transferred: {number}")
+
+
+@click.command(cls=ButlerCommand)
+@repo_argument(required=True)
+@click.argument("parent", required=True, nargs=1)
+@click.argument("children", required=False, nargs=-1)
+@click.option("--doc", default="",
+              help="Documentation string associated with this collection. "
+              "Only relevant if the collection is newly created.")
+@click.option("--flatten/--no-flatten", default=False,
+              help="If `True` recursively flatten out any nested chained collections in children first.")
+@click.option("--mode",
+              type=click.Choice(["redefine", "extend", "remove", "prepend", "pop"]),
+              default="redefine",
+              help="Update mode: "
+              "'redefine': Create new chain or redefine existing chain with the supplied CHILDREN. "
+              "'remove': Modify existing chain to remove the supplied CHILDREN. "
+              "'pop': Pop a numbered element off the chain. Defaults to popping "
+              "the first element (0). ``children`` must be integers if given. "
+              "'prepend': Modify existing chain to prepend the supplied CHILDREN to the front. "
+              "'extend': Modify existing chain to extend it with the supplied CHILDREN.")
+def collection_chain(**kwargs):
+    """Define a collection chain.
+
+    PARENT is the name of the chained collection to create or modify. If the
+    collection already exists the chain associated with it will be updated.
+
+    CHILDREN are the collections to be used to modify the chain. The exact
+    usage depends on the MODE option. When the MODE is 'pop' the CHILDREN
+    should be integer indices indicating collections to be removed from
+    the current chain.
+
+    MODE 'pop' can take negative integers to indicate removal relative to the
+    end of the chain, but when doing that '--' must be given to indicate the
+    end of the options specification.
+
+    $ butler collection-chain REPO --mode=pop PARENT -- -1
+
+    Will remove the final collection from the chain.
+    """
+    chain = script.collectionChain(**kwargs)
+    print(f"[{', '.join(chain)}]")

python/lsst/daf/butler/script/__init__.py

@@ -22,6 +22,7 @@
 from .butlerImport import butlerImport
 from .certifyCalibrations import certifyCalibrations
 from .createRepo import createRepo
+from .collectionChain import collectionChain
 from .configDump import configDump
 from .configValidate import configValidate
 from .pruneCollection import pruneCollection

python/lsst/daf/butler/script/collectionChain.py

@@ -0,0 +1,127 @@
+# This file is part of daf_butler.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (http://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from __future__ import annotations
+
+from .. import Butler, CollectionType
+from ..registry import MissingCollectionError
+
+
+def collectionChain(repo, mode, parent, children, doc, flatten):
+    """Get the collections whose names match an expression.
+
+    Parameters
+    ----------
+    repo : `str`
+        URI to the location of the repo or URI to a config file describing the
+        repo and its location.
+    mode : `str`
+        Update mode for this chain. Options are:
+        'redefine': Create or modify ``parent`` to be defined by the supplied
+        ``children``.
+        'remove': Modify existing chain to remove ``children`` from it.
+        'prepend': Add the given ``children`` to the beginning of the chain.
+        'extend': Modify existing chain to add ``children`` to the end of it.
+        'pop': Pop a numbered element off the chain. Defaults to popping
+        the first element (0). ``children`` must be integers if given.
+        Both 'prepend' and 'extend' are the same as 'redefine' if the chain
+        does not exist.
+    parent: `str`
+        Name of the chained collection to update. Will be created if it
+        does not exist already.
+    children: iterable of `str`
+        Names of the children to be included in the chain.
+    doc : `str`
+        If the chained collection is being created, the documentation string
+        that will be associated with it.
+    flatten : `str`
+        If `True`, recursively flatten out any nested
+        `~CollectionType.CHAINED` collections in ``children`` first.
+
+    Returns
+    -------
+    chain : `tuple` of `str`
+        The collections in the chain following this command.
+    """
+    butler = Butler(repo, writeable=True)
+
+    # Every mode needs children except pop.
+    if not children and mode != "pop":
+        raise RuntimeError(f"Must provide children when defining a collection chain in mode {mode}.")
+
+    try:
+        butler.registry.getCollectionType(parent)
+    except MissingCollectionError:
+        # Create it -- but only if mode can work with empty chain.
+        if mode in ("redefine", "extend", "prepend"):
+            if not doc:
+                doc = None
+            butler.registry.registerCollection(parent, CollectionType.CHAINED, doc)
+        else:
+            raise RuntimeError(f"Mode '{mode}' requires that the collection exists "
+                               f"but collection '{parent}' is not known to this registry") from None
+
+    current = list(butler.registry.getCollectionChain(parent))
+
+    if mode == "redefine":
+        # Given children are what we want.
+        pass
+    elif mode == "prepend":
+        children = tuple(children) + tuple(current)
+    elif mode == "extend":
+        current.extend(children)
+        children = current
+    elif mode == "remove":
+        for child in children:
+            current.remove(child)
+        children = current
+    elif mode == "pop":
+        if children:
+            n_current = len(current)
+
+            def convert_index(i):
+                """Convert negative index to positive."""
+                if i >= 0:
+                    return i
+                return n_current + i
+
+            # For this mode the children should be integers.
+            # Convert negative integers to positive ones to allow
+            # sorting.
+            children = [convert_index(int(child)) for child in children]
+
+            # Reverse sort order so we can remove from the end first
+            children = reversed(sorted(children))
+
+        else:
+            # Nothing specified, pop from the front of the chin.
+            children = [0]
+
+        for i in children:
+            current.pop(i)
+
+        children = current
+    else:
+        raise ValueError(f"Unrecognized update mode: '{mode}'")
+
+    butler.registry.setCollectionChain(parent, children, flatten=flatten)
+
+    return tuple(butler.registry.getCollectionChain(parent))