Expand the mode option and remove --pop

Pop is now a mode that allows integer indices to specify the
elements to pop.

Add extend option.

Also now prints the current contents of the chain on completion.

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

@@ -503,19 +503,20 @@ def transfer_datasets(**kwargs):
 @repo_argument(required=True)
 @click.argument("parent", required=True, nargs=1)
 @click.argument("children", required=False, nargs=-1)
-@click.option("--pop", is_flag=True, default=False,
-              help="Pop the first collection off the chain. Can not be used if CHILDREN are given.")
 @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"]),
+              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.
@@ -524,6 +525,17 @@ def collection_chain(**kwargs):
     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.
+    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.
     """
-    script.collectionChain(**kwargs)
+    chain = script.collectionChain(**kwargs)
+    print(f"[{', '.join(chain)}]")

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

@@ -25,7 +25,7 @@
 from ..registry import MissingCollectionError
 
 
-def collectionChain(repo, mode, parent, children, doc, flatten, pop):
+def collectionChain(repo, mode, parent, children, doc, flatten):
     """Get the collections whose names match an expression.
 
     Parameters
@@ -38,8 +38,12 @@ def collectionChain(repo, mode, parent, children, doc, flatten, pop):
         '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.
-        This is the same as 'redefine' if the chain does not exist.
+        '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.
@@ -51,29 +55,23 @@ def collectionChain(repo, mode, parent, children, doc, flatten, pop):
     flatten : `str`
         If `True`, recursively flatten out any nested
         `~CollectionType.CHAINED` collections in ``children`` first.
-    pop : `bool`
-        Pop the first collection off the chain. Can not be used if ``children``
-        are given. The ``parent`` collection must exist.
+
+    Returns
+    -------
+    chain : `tuple` of `str`
+        The collections in the chain following this command.
     """
     butler = Butler(repo, writeable=True)
 
-    if pop:
-        if children:
-            raise RuntimeError("Can not provide children if using 'pop'")
-
-        children = butler.registry.getCollectionChain(parent)[1:]
-        butler.registry.setCollectionChain(parent, children)
-        return
-
-    # All other options require children.
-    if not children:
-        raise RuntimeError("Must provide children when defining a collection chain.")
+    # 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 is redefine or extend
-        if mode in ("redefine", "extend"):
+        # 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)
@@ -86,14 +84,43 @@ def collectionChain(repo, mode, parent, children, doc, flatten, pop):
     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:
+            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))