Add butler collection-chain command

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,22 @@ 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=True, 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.")
+def collection_chain(**kwargs):
+    """Define a collection chain.
+
+    PARENT is the name of the chained collection to create. If the collection
+    already exists the chain associated with it will be updated.
+    CHILDREN are the collections to be included in the chain in order.
+    """
+    script.collectionChain(**kwargs)

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,59 @@
+# 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, 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.
+    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.
+    """
+
+    butler = Butler(repo, writeable=True)
+
+    try:
+        butler.registry.getCollectionType(parent)
+    except MissingCollectionError:
+        # Create it
+        if not doc:
+            doc = None
+        butler.registry.registerCollection(parent, CollectionType.CHAINED, doc)
+
+    butler.registry.setCollectionChain(parent, children, flatten=flatten)

tests/test_cliCmdQueryCollections.py

@@ -135,10 +135,12 @@ def testChained(self):
             registry1.registerRun("run1")
             registry1.registerCollection("tag1", CollectionType.TAGGED)
             registry1.registerCollection("calibration1", CollectionType.CALIBRATION)
-            registry1.registerCollection("chain1", CollectionType.CHAINED)
-            registry1.registerCollection("chain2", CollectionType.CHAINED)
-            registry1.setCollectionChain("chain1", ["tag1", "run1", "chain2"])
-            registry1.setCollectionChain("chain2", ["calibration1", "run1"])
+
+            # Create the collection chain
+            result = self.runner.invoke(cli, ["collection-chain", "here", "chain2", "calibration1", "run1"])
+            self.assertEqual(result.exit_code, 0, clickResultMsg(result))
+            result = self.runner.invoke(cli, ["collection-chain", "here", "chain1", "tag1", "run1", "chain2"])
+            self.assertEqual(result.exit_code, 0, clickResultMsg(result))
 
             # Use the script function to test the query-collections TREE
             # option, because the astropy.table.Table.read method, which we are
@@ -154,15 +156,15 @@ def testChained(self):
                                     ("run1", "RUN"),
                                     ("tag1", "TAGGED"),
                                     ("calibration1", "CALIBRATION"),
+                                    ("chain2", "CHAINED"),
+                                    ("  calibration1", "CALIBRATION"),
+                                    ("  run1", "RUN"),
                                     ("chain1", "CHAINED"),
                                     ("  tag1", "TAGGED"),
                                     ("  run1", "RUN"),
                                     ("  chain2", "CHAINED"),
                                     ("    calibration1", "CALIBRATION"),
-                                    ("    run1", "RUN"),
-                                    ("chain2", "CHAINED"),
-                                    ("  calibration1", "CALIBRATION"),
-                                    ("  run1", "RUN"))),
+                                    ("    run1", "RUN"))),
                              names=("Name", "Type"))
             self.assertAstropyTablesEqual(table, expected)
 
@@ -174,8 +176,8 @@ def testChained(self):
                 ("run1", "RUN", ""),
                 ("tag1", "TAGGED", ""),
                 ("calibration1", "CALIBRATION", ""),
-                ("chain1", "CHAINED", "[tag1, run1, chain2]"),
-                ("chain2", "CHAINED", "[calibration1, run1]"))),
+                ("chain2", "CHAINED", "[calibration1, run1]"),
+                ("chain1", "CHAINED", "[tag1, run1, chain2]"))),
                 names=("Name", "Type", "Definition"))
             table = readTable(result.output)
             self.assertAstropyTablesEqual(readTable(result.output), expected)
@@ -188,11 +190,11 @@ def testChained(self):
                 ("run1", "RUN"),
                 ("tag1", "TAGGED"),
                 ("calibration1", "CALIBRATION"),
+                ("calibration1", "CALIBRATION"),
+                ("run1", "RUN"),
                 ("tag1", "TAGGED"),
                 ("run1", "RUN"),
-                ("calibration1", "CALIBRATION"),
-                ("calibration1", "CALIBRATION"),
-                ("run1", "RUN"))),
+                ("calibration1", "CALIBRATION"))),
                 names=("Name", "Type"))
             self.assertAstropyTablesEqual(readTable(result.output), expected)