extract information from STAC extensions

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"
+    }
+  ]
+}

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

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

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

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"]

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)

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]