diff --git a/src/stactools/cli/__init__.py b/src/stactools/cli/__init__.py index 7416e760..7b3e74fd 100644 --- a/src/stactools/cli/__init__.py +++ b/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) diff --git a/src/stactools/cli/commands/add.py b/src/stactools/cli/commands/add.py new file mode 100644 index 00000000..70c3dd3b --- /dev/null +++ b/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 diff --git a/src/stactools/core/__init__.py b/src/stactools/core/__init__.py index e7686be3..b823298e 100644 --- a/src/stactools/core/__init__.py +++ b/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" diff --git a/src/stactools/core/add.py b/src/stactools/core/add.py new file mode 100644 index 00000000..ff14629e --- /dev/null +++ b/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." + ) diff --git a/tests/cli/commands/test_add.py b/tests/cli/commands/test_add.py new file mode 100644 index 00000000..9fd9d9bc --- /dev/null +++ b/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)