Source code for registrar.route.stac

"""Module holding classes dealing with registration routes of STAC objects
"""

import json
from typing import List, Optional

import pystac

from ..abc import Route
from ..config import SourceConfig
from ..source import Source, get_source


class ItemRoute(Route[pystac.Item]):
    """A route implementation for STAC Items"""

    def parse(self, raw: str) -> pystac.Item:
        return pystac.Item.from_dict(json.loads(raw))

    def get_source(
        self, source_cfgs: List[SourceConfig], item: pystac.Item
    ) -> Optional[Source]:
        return get_source(
            source_cfgs, [asset.href for asset in item.assets.values()]
        )


class CollectionRoute(Route[pystac.Catalog]):
    """A route implementation for STAC Collections"""

    def parse(self, raw: str) -> pystac.Catalog:
        values = json.loads(raw)
        if values.get("type") == "Collection":
            return pystac.Collection.from_dict(values)
        elif values.get("type") == "Catalog":
            return pystac.Catalog.from_dict(values)
        else:
            raise ValueError(f"Invalid collection/catalog object {values}")

    def get_source(
        self, source_cfgs: List[SourceConfig], item: pystac.Catalog
    ) -> Optional[Source]:
        if isinstance(item, pystac.Collection):
            return get_source(
                source_cfgs, [asset.href for asset in item.assets.values()]
            )
        return get_source(
            source_cfgs, [item.get_self_href() or ""]
        )