diff --git a/tests/data/alternate_asset.json b/tests/data/alternate_asset.json new file mode 100644 index 0000000..9ce1852 --- /dev/null +++ b/tests/data/alternate_asset.json @@ -0,0 +1,99 @@ +{ + "stac_version": "1.0.0", + "stac_extensions": [ + "https://stac-extensions.github.io/storage/v2.0.0/schema.json", + "https://stac-extensions.github.io/version/v1.2.0/schema.json", + "https://stac-extensions.github.io/alternate-assets/v1.1.0/schema.json" + ], + "id": "m_3009743_sw_14_1_20160928_20161129", + "bbox": [ + -97.75, + 30.25, + -97.6875, + 30.312499999999996 + ], + "type": "Feature", + "geometry": { + "type": "Polygon", + "coordinates": [ + [ + [ + -97.75, + 30.312499999999996 + ], + [ + -97.75, + 30.25 + ], + [ + -97.6875, + 30.25 + ], + [ + -97.6875, + 30.312499999999996 + ], + [ + -97.75, + 30.312499999999996 + ] + ] + ] + }, + "properties": { + "datetime": "2016-09-28T00:00:00+00:00", + "mission": "NAIP", + "gsd": 1, + "storage:schemes": { + "az-wus2-ar": { + "type": "ms-azure", + "platform": "https://{account}.blob.core.windows.net", + "account": "jon-doe-123", + "region": "westus2" + }, + "aws-std": { + "type": "aws-s3", + "platform": "https://{bucket}.s3.{region}.amazonaws.com", + "bucket": "naip-visualization", + "region": "us-west-2", + "requester_pays": true + }, + "minio": { + "type": "custom-s3", + "platform": "https://play.min.io:9000" + } + } + }, + "assets": { + "CO_GEOTIFF_RGB": { + "href": "s3://naip-visualization/tx/2016/100cm/rgb/30097/m_3009743_sw_14_1_20160928.tif", + "type": "image/tiff; application=geotiff; profile=cloud-optimized", + "storage:refs": [ + "aws-std" + ], + "alternate:name": "Amazon S3", + "alternate": { + "minio": { + "href": "s3://mybucket/tx/2016/100cm/rgb/30097/m_3009743_sw_14_1_20160928.tif", + "alternate:name": "MinIO", + "storage:refs": [ + "minio" + ] + } + } + }, + "GEOTIFF_AZURE_RGBIR": { + "href": "https://naip-nsl.blob.core.windows.net/tx/2016/100cm/rgb/30097/m_3009743_sw_14_1_20160928.tif", + "type": "image/tiff; application=geotiff", + "storage:refs": [ + "az-wus2-ar" + ] + } + }, + "links": [ + { + "href": "https://example.com/examples/item-naip.json", + "rel": "self" + } + ] +} diff --git a/tests/test_extensions.py b/tests/test_extensions.py new file mode 100644 index 0000000..a9973d8 --- /dev/null +++ b/tests/test_extensions.py @@ -0,0 +1,70 @@ +import json +import pathlib + +import pystac +import pytest + +from xpystac import extensions +from xpystac.extensions.common import JSON, AssetInfo + + +@pytest.fixture(scope="module") +def alternate_asset_item(request): + root = pathlib.Path(__file__).parent / "data" + path = root / "alternate_asset.json" + data = json.loads(path.read_text()) + + return pystac.Item.from_dict(data) + + +@pytest.mark.parametrize( + ["alternate", "expected_href", "expected_properties"], + ( + pytest.param( + None, + "s3://naip-visualization/tx/2016/100cm/rgb/30097/m_3009743_sw_14_1_20160928.tif", + {"alternate:name": "Amazon S3", "storage:refs": ["aws-std"]}, + id="default", + ), + pytest.param( + "minio", + "s3://mybucket/tx/2016/100cm/rgb/30097/m_3009743_sw_14_1_20160928.tif", + {"alternate:name": "MinIO", "storage:refs": ["minio"]}, + id="minio", + ), + ), +) +def test_extract_alternate_asset( + alternate_asset_item: pystac.Item, + alternate: str | None, + expected_href: str, + expected_properties: dict[str, JSON], +) -> None: + asset = alternate_asset_item.assets["CO_GEOTIFF_RGB"] + + actual: AssetInfo = extensions.extract_alternate_asset(asset, alternate) + + assert actual.href == expected_href + assert actual.properties == expected_properties + + +@pytest.mark.parametrize( + ["refs", "expected"], + ( + pytest.param([], [], id="none"), + pytest.param(["minio"], [{"type": "custom-s3"}], id="one"), + pytest.param( + ["minio", "aws-us-west-2"], + [{"type": "custom-s3"}, {"type": "aws-s3"}], + id="two", + ), + ), +) +def test_resolve_refs(refs: list[str], expected: list[dict[str, JSON]]) -> None: + schemes: dict[str, dict[str, JSON]] = { + "minio": {"type": "custom-s3"}, + "aws-us-west-2": {"type": "aws-s3"}, + } + + actual = extensions.schemes._resolve_refs(refs, schemes) + assert actual == expected diff --git a/xpystac/core.py b/xpystac/core.py index 6dc2f49..d6dc75c 100644 --- a/xpystac/core.py +++ b/xpystac/core.py @@ -1,10 +1,13 @@ import functools from collections.abc import Callable +from typing import cast import pystac import xarray +from xpystac import extensions from xpystac._xstac_kerchunk import _stac_to_kerchunk +from xpystac.extensions import JSON from xpystac.utils import _import_optional_dependency @@ -83,11 +86,26 @@ def _( obj: pystac.Asset, patch_url: None | Callable[[str], str] = None, allow_kerchunk: bool = True, + *, + alternate: str | None = None, **kwargs, ) -> xarray.Dataset: - open_kwargs = obj.extra_fields.get("xarray:open_kwargs", {}) + # extract extensions + info = extensions._extract_alternate_asset(obj, alternate) + href = info.href + + storage = extensions.extract_scheme(info, kind="storage") + auth = extensions.extract_scheme(info, kind="auth") - storage_options = obj.extra_fields.get("xarray:storage_options", None) + open_kwargs = cast(dict[str, JSON], info.properties.get("xarray:open_kwargs", {})) + storage_options = cast( + dict[str, JSON] | None, info.properties.get("xarray:storage_options", None) + ) + + # use extension data + # FIXME: these should be replaced by the authentication, zarr and storage extensions + if patch_url is not None: + href = patch_url(href) if storage_options: open_kwargs["storage_options"] = storage_options @@ -97,7 +115,7 @@ def _( and {"index", "references"}.intersection(set(obj.roles) if obj.roles else set()) ): requests = _import_optional_dependency("requests") - r = requests.get(obj.href) + r = requests.get(href) r.raise_for_status() refs = r.json() @@ -125,9 +143,5 @@ def _( return read_icechunk(obj) - href = obj.href - if patch_url is not None: - href = patch_url(href) - ds = xarray.open_dataset(href, **{**default_kwargs, **open_kwargs, **kwargs}) return ds diff --git a/xpystac/extensions.py b/xpystac/extensions.py new file mode 100644 index 0000000..53e69d4 --- /dev/null +++ b/xpystac/extensions.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +JSON = dict[str, "JSON"] | list["JSON"] | str | int | float | bool | None + +if TYPE_CHECKING: + from typing import Literal + + import pystac + + +@dataclass +class AssetInfo: + href: str + properties: dict[str, JSON] + + +def _extract_alternate_asset(asset: pystac.Asset, alternate: str | None) -> AssetInfo: + excludes = {"href", "alternate"} + + properties = {k: v for k, v in asset.extra_fields.items() if k not in excludes} + if alternate is None: + href = asset.href + additional_properties = {} + else: + alternate_assets = asset.extra_fields.get("alternate") + if alternate_assets is None: + raise ValueError( + "Alternate asset name given but no alternate assets found." + ) + + alternate_asset = alternate_assets.get(alternate) + if alternate_asset is None: + raise ValueError( + f"alternate asset {alternate} not found." + " Choose one of {', '.join(sorted(alternate_assets))}" + ) + + href = alternate_asset["href"] + additional_properties = { + k: v for k, v in alternate_asset.items() if k not in excludes + } + + return AssetInfo(href, properties | additional_properties) + + +def _extract_parent_attribute(obj: pystac.Asset, attr: str) -> JSON: + if isinstance(obj.owner, pystac.Item): + fields = obj.owner.properties + else: + fields = obj.owner.extra_fields + + return fields.get(attr) + + +def _resolve_refs( + refs: list[str], schemes: dict[str, dict[str, JSON]] +) -> list[dict[str, JSON]]: + missing_refs = [ref for ref in refs if ref not in schemes] + if missing_refs: + raise ValueError("selected unknown refs: {', '.join(missing_refs)}") + + return [schemes[ref] for ref in refs] + + +def extract_scheme( + info: AssetInfo, kind: Literal["storage", "auth"] +) -> dict[str, JSON] | None: + refs = cast(list[str], info.properties.get(f"{kind}:refs", [])) + if not refs: + return None + + schemes = cast( + dict[str, dict[str, JSON]], _extract_parent_attribute(info, f"{kind}:schemes") + ) + if schemes is None: + raise ValueError( + f"{kind}:refs found but no {kind}:schemes on the parent object" + ) + + if len(refs) != 1: + raise NotImplementedError( + f"Only one {kind}:ref per asset is currently supported" + ) + [scheme] = _resolve_refs(refs, schemes) + + return scheme diff --git a/xpystac/extensions/__init__.py b/xpystac/extensions/__init__.py new file mode 100644 index 0000000..d431d33 --- /dev/null +++ b/xpystac/extensions/__init__.py @@ -0,0 +1,4 @@ +from xpystac.extensions.alternate import extract_alternate_asset +from xpystac.extensions.schemes import extract_scheme + +__all__ = ["extract_alternate_asset", "extract_scheme"] diff --git a/xpystac/extensions/alternate.py b/xpystac/extensions/alternate.py new file mode 100644 index 0000000..1cdbe9a --- /dev/null +++ b/xpystac/extensions/alternate.py @@ -0,0 +1,32 @@ +import pystac + +from xpystac.extensions.common import AssetInfo + + +def extract_alternate_asset(asset: pystac.Asset, alternate: str | None) -> AssetInfo: + excludes = {"href", "alternate"} + + properties = {k: v for k, v in asset.extra_fields.items() if k not in excludes} + if alternate is None: + href = asset.href + additional_properties = {} + else: + alternate_assets = asset.extra_fields.get("alternate") + if alternate_assets is None: + raise ValueError( + "Alternate asset name given but no alternate assets found." + ) + + alternate_asset = alternate_assets.get(alternate) + if alternate_asset is None: + raise ValueError( + f"alternate asset {alternate} not found." + " Choose one of {', '.join(sorted(alternate_assets))}" + ) + + href = alternate_asset["href"] + additional_properties = { + k: v for k, v in alternate_asset.items() if k not in excludes + } + + return AssetInfo(href, properties | additional_properties) diff --git a/xpystac/extensions/common.py b/xpystac/extensions/common.py new file mode 100644 index 0000000..97d4969 --- /dev/null +++ b/xpystac/extensions/common.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + +from xpystac.typing import JSON + + +@dataclass +class AssetInfo: + href: str + properties: dict[str, JSON] diff --git a/xpystac/extensions/schemes.py b/xpystac/extensions/schemes.py new file mode 100644 index 0000000..9f385a9 --- /dev/null +++ b/xpystac/extensions/schemes.py @@ -0,0 +1,48 @@ +from typing import Literal, cast + +import pystac + +from xpystac.typing import JSON + + +def _resolve_refs( + refs: list[str], schemes: dict[str, dict[str, JSON]] +) -> list[dict[str, JSON]]: + missing_refs = [ref for ref in refs if ref not in schemes] + if missing_refs: + raise ValueError("selected unknown refs: {', '.join(missing_refs)}") + + return [schemes[ref] for ref in refs] + + +def _extract_parent_attribute(obj: pystac.Asset, attr: str) -> JSON: + if isinstance(obj.owner, pystac.Item): + fields = obj.owner.properties + else: + fields = obj.owner.extra_fields + + return fields.get(attr) + + +def extract_scheme( + asset: pystac.Asset, kind: Literal["storage", "auth"] +) -> dict[str, JSON] | None: + refs = cast(list[str], asset.properties.get(f"{kind}:refs", [])) + if not refs: + return None + + schemes = cast( + dict[str, dict[str, JSON]], _extract_parent_attribute(asset, f"{kind}:schemes") + ) + if schemes is None: + raise ValueError( + f"{kind}:refs found but no {kind}:schemes on the parent object" + ) + + if len(refs) != 1: + raise NotImplementedError( + f"Only one {kind}:ref per asset is currently supported" + ) + [scheme] = _resolve_refs(refs, schemes) + + return scheme diff --git a/xpystac/typing.py b/xpystac/typing.py new file mode 100644 index 0000000..9e4f7e2 --- /dev/null +++ b/xpystac/typing.py @@ -0,0 +1,3 @@ +from __future__ import annotations + +JSON = dict[str, "JSON"] | list["JSON"] | str | int | float | bool | None