Added 'add' command to add an item to a catalog/collection (#153)

Co-authored-by: Pete Gadomski <pete.gadomski@gmail.com>
Co-authored-by: Tyler Battle <tyler@battle.ca>

src/stactools/cli/__init__.py

@@ -8,8 +8,8 @@
 def register_plugin(registry: 'Registry') -> None:
     # Register subcommands
 
-    from stactools.cli.commands import (copy, info, layout, merge, migrate,
-                                        version, validate)
+    from stactools.cli.commands import (add, copy, info, layout, merge,
+                                        migrate, validate, version)
 
     registry.register_subcommand(copy.create_copy_command)
     registry.register_subcommand(copy.create_move_assets_command)
@@ -19,6 +19,7 @@ def register_plugin(registry: 'Registry') -> None:
     registry.register_subcommand(merge.create_merge_command)
     registry.register_subcommand(validate.create_validate_command)
     registry.register_subcommand(version.create_version_command)
+    registry.register_subcommand(add.create_add_command)
 
     # TODO
     # registry.register_subcommand(migrate.create_migrate_command)

src/stactools/cli/commands/add.py

@@ -0,0 +1,54 @@
+import click
+
+from typing import Optional
+
+from pystac import Catalog, Item, read_file
+
+from stactools.core import add_item
+
+
+def add(source_item: str,
+        target_catalog: str,
+        collection_id: Optional[str] = None,
+        move_assets: bool = False) -> None:
+    source = read_file(source_item)
+    if not isinstance(source, Item):
+        raise click.BadArgumentUsage(f"{source_item} is not a STAC Item")
+    target = read_file(target_catalog)
+    if not isinstance(target, Catalog):
+        raise click.BadArgumentUsage(f"{target_catalog} is not a STAC Catalog")
+
+    if collection_id is not None:
+        target_collection = target.get_child(collection_id, recursive=True)
+        if target_collection is None:
+            raise click.BadOptionUsage(
+                'collection',
+                'A collection with ID {} does not exist in {}'.format(
+                    collection_id, target_catalog))
+
+        add_item(source, target_collection, move_assets)
+        target_collection.save()
+    else:
+        add_item(source, target, move_assets)
+        target.save()
+
+
+def create_add_command(cli: click.Group) -> click.Command:
+    @cli.command('add', short_help='Add an item to a catalog/collection.')
+    @click.argument('source_item')
+    @click.argument('target_catalog')
+    @click.option('--collection',
+                  help=("The collection ID to add to. If not set, will "
+                        "add to the root catalog or collection."))
+    @click.option('-a',
+                  '--move-assets',
+                  is_flag=True,
+                  help='Move assets to the target catalog Item locations.')
+    def add_command(source_item: str, target_catalog: str, collection: str,
+                    move_assets: bool) -> None:
+        add(source_item,
+            target_catalog,
+            collection_id=collection,
+            move_assets=move_assets)
+
+    return add_command

src/stactools/core/__init__.py

@@ -5,5 +5,6 @@
                                  move_all_assets, copy_catalog)
 from stactools.core.layout import layout_catalog
 from stactools.core.merge import (merge_items, merge_all_items)
+from stactools.core.add import add_item
 
 __version__ = "0.2.2"

src/stactools/core/add.py

@@ -0,0 +1,47 @@
+import os
+
+from pystac import Catalog, Item, Collection
+from pystac.layout import BestPracticesLayoutStrategy
+
+from stactools.core.copy import move_assets as do_move_assets
+
+
+def add_item(source_item: Item,
+             target_catalog: Catalog,
+             move_assets: bool = False) -> None:
+    """Add a item into a catalog.
+
+    Args:
+        source_item (pystac.Item): The Item that will be added.
+            This item is not mutated in this operation.
+        target_catalog (pystac.Item): The destination catalog.
+            This catalog will be mutated in this operation.
+        move_assets (bool): If true, move the asset files alongside the target item.
+    """
+
+    target_item_ids = [item.id for item in target_catalog.get_all_items()]
+    if source_item.id in target_item_ids:
+        raise ValueError(
+            f'An item with ID {source_item.id} already exists in the target catalog'
+        )
+    self_href = target_catalog.get_self_href()
+    if self_href:
+        parent_dir = os.path.dirname(self_href)
+        layout_strategy = BestPracticesLayoutStrategy()
+        item_copy = source_item.clone()
+        item_copy.set_self_href(
+            layout_strategy.get_item_href(item_copy, parent_dir))
+        target_catalog.add_item(item_copy)
+
+        if isinstance(target_catalog, Collection):
+            item_copy.set_collection(target_catalog)
+            target_catalog.update_extent_from_items()
+        else:
+            item_copy.set_collection(None)
+
+        if move_assets:
+            do_move_assets(item_copy, copy=False)
+    else:
+        raise ValueError(
+            f"Cannot add Item {source_item.id} because {target_catalog} does not have a self href."
+        )

tests/cli/commands/test_add.py

@@ -0,0 +1,91 @@
+from tempfile import TemporaryDirectory
+
+import pystac
+from stactools.core import move_all_assets
+from stactools.cli.commands.add import create_add_command
+from stactools.testing import CliTestCase
+from .test_cases import TestCases
+
+
+def create_temp_catalog_copy(tmp_dir):
+    col = TestCases.planet_disaster()
+    col.normalize_hrefs(tmp_dir)
+    col.save(catalog_type=pystac.CatalogType.SELF_CONTAINED)
+    move_all_assets(col, copy=True)
+    col.save()
+
+    return col
+
+
+class AddTest(CliTestCase):
+    def create_subcommand_functions(self):
+        return [create_add_command]
+
+    def test_add_item(self):
+        catalog = TestCases.test_case_1()
+        subcatalog = list(list(catalog.get_children())[0].get_children())[0]
+        item = list(subcatalog.get_all_items())[0]
+        item_path = item.get_self_href()
+        with TemporaryDirectory() as tmp_dir:
+            target_catalog = create_temp_catalog_copy(tmp_dir)
+
+            items = list(target_catalog.get_all_items())
+            self.assertEqual(len(items), 5)
+
+            cmd = ["add", item_path, target_catalog.get_self_href()]
+
+            self.run_command(cmd)
+
+            target_col = pystac.read_file(target_catalog.get_self_href())
+            items = list(target_col.get_all_items())
+            self.assertEqual(len(items), 6)
+
+    def test_add_item_to_specific_collection(self):
+        catalog = TestCases.test_case_1()
+        subcatalog = list(list(catalog.get_children())[0].get_children())[0]
+        item = list(subcatalog.get_all_items())[0]
+        item_path = item.get_self_href()
+        with TemporaryDirectory() as tmp_dir:
+            target_catalog = create_temp_catalog_copy(tmp_dir)
+            items = list(target_catalog.get_all_items())
+            self.assertEqual(len(items), 5)
+
+            cmd = [
+                "add",
+                item_path,
+                target_catalog.get_self_href(),
+                "--collection",
+                "hurricane-harvey",
+            ]
+
+            res = self.run_command(cmd)
+            self.assertEqual(res.exit_code, 0)
+
+            target_col = pystac.read_file(target_catalog.get_self_href())
+            child_col = target_col.get_child("hurricane-harvey")
+            target_item = child_col.get_item(item.id)
+            self.assertIsNotNone(target_item)
+
+    def test_add_item_to_missing_collection(self):
+        catalog = TestCases.test_case_1()
+        subcatalog = list(list(catalog.get_children())[0].get_children())[0]
+        item = list(subcatalog.get_all_items())[0]
+        item_path = item.get_self_href()
+        with TemporaryDirectory() as tmp_dir:
+            target_catalog = create_temp_catalog_copy(tmp_dir)
+
+            items = list(target_catalog.get_all_items())
+            self.assertEqual(len(items), 5)
+
+            cmd = [
+                "add",
+                item_path,
+                target_catalog.get_self_href(),
+                "--collection",
+                "WRONG",
+            ]
+
+            res = self.run_command(cmd)
+            self.assertEqual(res.exit_code, 2)
+            self.assertTrue(
+                " A collection with ID WRONG does not exist" in res.output)