add butler associate command

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

@@ -19,7 +19,8 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-__all__ = ("butler_import",
+__all__ = ("associate",
+           "butler_import",
            "certify_calibrations",
            "create",
            "config_dump",
@@ -32,10 +33,11 @@
            "query_datasets",
            "query_dimension_records",
            "remove_dataset_type",
-           )
+)
 
 
-from .commands import (butler_import,
+from .commands import (associate,
+                       butler_import,
                        certify_calibrations,
                        create,
                        config_dump,
@@ -48,4 +50,4 @@
                        query_datasets,
                        query_dimension_records,
                        remove_dataset_type,
-                       )
+)

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

@@ -60,6 +60,18 @@
 existingRepoHelp = "REPO is the URI or path to an existing data repository root or configuration file."
 
 
+@click.command(cls=ButlerCommand, short_help="Add existing datasets to a tagged collection.")
+@repo_argument(required=True)
+@collection_argument(help="COLLECTION is the collection the datasets should be associated with.")
+@query_datasets_options(repo=False, showUri=False, useArguments=False)
+@options_file_option()
+def associate(**kwargs):
+    """Add existing datasets to a tagged collection; searches for datasets with
+    the options and adds them to the named COLLECTION.
+    """
+    script.associate(**kwargs)
+
+
 # The conversion from the import command name to the butler_import function
 # name for subcommand lookup is implemented in the cli/butler.py, in
 # funcNameToCmdName and cmdNameToFuncName. If name changes are made here they

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

@@ -28,6 +28,7 @@
 from .queryCollections import queryCollections
 from .queryDataIds import queryDataIds
 from .queryDatasets import QueryDatasets
+from ._associate import associate # depends on QueryDatasets
 from ._pruneDatasets import pruneDatasets # depends on QueryDatasets
 from .queryDatasetTypes import queryDatasetTypes
 from .queryDimensionRecords import queryDimensionRecords

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

@@ -0,0 +1,45 @@
+# 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 .. import Butler, CollectionType
+from ..script import QueryDatasets
+
+
+def associate(repo, collection, dataset_type, collections, where, find_first):
+    """Add existing datasets to ta CHAINED collection.
+    """
+
+    butler = Butler(repo, writeable=True)
+
+    if butler.registry.getCollectionType(collection) is not CollectionType.CHAINED:
+        raise RuntimeError("Can only add datasets to a chained collection.")
+
+    results = QueryDatasets(
+        butler=butler,
+        glob=dataset_type,
+        collections=collections,
+        where=where,
+        find_first=find_first,
+        show_uri=False,
+        repo=None
+    )
+
+    butler.registry.associate(collection, results.getDatasets())

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

@@ -110,9 +110,13 @@ class QueryDatasets:
 
     Parameters
     ----------
-    repo : `str`
+    repo : `str` or `None`
         URI to the location of the repo or URI to a config file describing the
-        repo and its location.
+        repo and its location. One of `repo` and `butler` must be `None` and
+        the other must not be `None`.
+    butler : ``lsst.daf.butler.Butler`` or `None`
+        The butler to use to query. One of `repo` and `butler` must be `None`
+        and the other must not be `None`.
     glob : iterable [`str`]
         A list of glob-style search string that fully or partially identify
         the dataset type names to search for.
@@ -133,8 +137,10 @@ class QueryDatasets:
         If True, include the dataset URI in the output.
     """
 
-    def __init__(self, repo, glob, collections, where, find_first, show_uri):
-        self.butler = Butler(repo)
+    def __init__(self, glob, collections, where, find_first, show_uri, repo=None, butler=None):
+        if (repo and butler) or (not repo and not butler):
+            raise RuntimeError("One of repo and butler must be provided and the other must be None.")
+        self.butler = butler or Butler(repo)
         self._getDatasets(glob, collections, where, find_first)
         self.showUri = show_uri
 

tests/test_cliCmdAssociate.py

@@ -0,0 +1,81 @@
+# 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/>.
+
+"""Unit tests for daf_butler CLI prune-datasets subcommand.
+"""
+
+import unittest
+from unittest.mock import patch
+
+from lsst.daf.butler.cli.butler import cli as butlerCli
+from lsst.daf.butler.cli.utils import clickResultMsg, LogCliRunner
+
+
+class AssociateTestCase(unittest.TestCase):
+    """Tests the ``associate`` ``butler`` subcommand.
+
+    ``script.associate`` contains no logic, so instead of mocking the
+    internals, just mock the call to that function to test for expected inputs
+    and input types.
+    """
+
+    def setUp(self):
+        self.runner = LogCliRunner()
+
+    @patch("lsst.daf.butler.script.associate")
+    def test_defaults(self, mockAssociate):
+        """Test the expected default values & types for optional options.
+        """
+        result = self.runner.invoke(
+            butlerCli, ["associate", "myRepo", "myCollection"])
+        self.assertEqual(result.exit_code, 0, clickResultMsg(result))
+        mockAssociate.assert_called_once_with(
+            repo="myRepo",
+            collection="myCollection",
+            dataset_type=tuple(),
+            collections=tuple(),
+            where=None,
+            find_first=False
+        )
+
+    @patch("lsst.daf.butler.script.associate")
+    def test_values(self, mockAssociate):
+        """Test expected values & types when passing in options.
+        """
+        result = self.runner.invoke(
+            butlerCli, ["associate", "myRepo", "myCollection",
+                        "--dataset-type", "myDatasetType",
+                        "--collections", "myCollection,otherCollection",
+                        "--where", "'a=b'",
+                        "--find-first"])
+        self.assertEqual(result.exit_code, 0, clickResultMsg(result))
+        mockAssociate.assert_called_once_with(
+            repo="myRepo",
+            collection="myCollection",
+            dataset_type=("myDatasetType",),
+            collections=("myCollection", "otherCollection"),
+            where="'a=b'",
+            find_first=True
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()