Xarray Assets Extension class (#1161)

* Add XarrayAssetsExtension class * Add tests and docs * Add some `open` tests * Docstring for open * Add xpystac to test deps * Update pystac/extensions/xarray_assets.py Co-authored-by: Pete Gadomski <[email protected]> * Remove open method * Update changelog * Fix copypasta --------- Co-authored-by: Pete Gadomski <[email protected]>
stac-utils · Jun 22, 2023 · f9e5bd6 · f9e5bd6
1 parent dd999cd
commit f9e5bd6
Show file tree

Hide file tree

Showing 9 changed files with 445 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,6 +20,7 @@
 - `KML` as a built in media type ([#1127](https://github.com/stac-utils/pystac/issues/1127))
 - `move/copy/delete` operations for local Assets ([#1158](https://github.com/stac-utils/pystac/issues/1158))
 - Latest core STAC spec jsonshemas are included in pytstac and used for validation ([#1165](https://github.com/stac-utils/pystac/pull/1165))
+- Xarray Assets Extension class ([#1161](https://github.com/stac-utils/pystac/pull/1161))
 
 ### Changed
 

diff --git a/docs/api.rst b/docs/api.rst
@@ -107,6 +107,7 @@ PySTAC provides support for the following STAC Extensions:
 * :mod:`Timestamps <pystac.extensions.timestamps>`
 * :mod:`Versioning Indicators <pystac.extensions.version>`
 * :mod:`View Geometry <pystac.extensions.view>`
+* :mod:`Xarray Assets <pystac.extensions.xarray_assets>`
 
 The following classes are used internally to implement these extensions and may be used
 to create custom implementations of STAC Extensions not supported by the library (see

diff --git a/docs/api/extensions.rst b/docs/api/extensions.rst
@@ -32,3 +32,4 @@ pystac.extensions
     timestamps.TimestampsExtension
     version.VersionExtension
     view.ViewExtension
+    xarray_assets.XarrayAssetsExtension
diff --git a/docs/api/extensions/xarray_assets.rst b/docs/api/extensions/xarray_assets.rst
@@ -0,0 +1,6 @@
+pystac.extensions.xarray_assets
+===============================
+
+.. automodule:: pystac.extensions.xarray_assets
+   :members:
+   :undoc-members:
diff --git a/pystac/__init__.py b/pystac/__init__.py
@@ -104,6 +104,7 @@
 import pystac.extensions.timestamps
 import pystac.extensions.version
 import pystac.extensions.view
+import pystac.extensions.xarray_assets
 
 EXTENSION_HOOKS = pystac.extensions.hooks.RegisteredExtensionHooks(
     [
@@ -126,6 +127,7 @@
         pystac.extensions.timestamps.TIMESTAMPS_EXTENSION_HOOKS,
         pystac.extensions.version.VERSION_EXTENSION_HOOKS,
         pystac.extensions.view.VIEW_EXTENSION_HOOKS,
+        pystac.extensions.xarray_assets.XARRAY_ASSETS_EXTENSION_HOOKS,
     ]
 )
 

diff --git a/pystac/extensions/xarray_assets.py b/pystac/extensions/xarray_assets.py
@@ -0,0 +1,162 @@
+"""Implements the :stac-ext:`Xarray Assets Extension <xarray>`."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Generic, List, Optional, TypeVar, Union
+
+import pystac
+from pystac.extensions.base import ExtensionManagementMixin, PropertiesExtension
+from pystac.extensions.hooks import ExtensionHooks
+
+T = TypeVar("T", pystac.Collection, pystac.Item, pystac.Asset)
+
+SCHEMA_URI = "https://stac-extensions.github.io/xarray-assets/v1.0.0/schema.json"
+
+PREFIX: str = "xarray:"
+OPEN_KWARGS_PROP = PREFIX + "open_kwargs"
+STORAGE_OPTIONS_PROP = PREFIX + "storage_options"
+
+
+class XarrayAssetsExtension(
+    Generic[T],
+    PropertiesExtension,
+    ExtensionManagementMixin[Union[pystac.Collection, pystac.Item]],
+):
+    """An abstract class that can be used to extend the properties of a
+    :class:`~pystac.Collection`, :class:`~pystac.Item`, or :class:`~pystac.Asset` with
+    properties from the :stac-ext:`Xarray Assets Extension <xarray>`. This class is
+    generic over the type of STAC Object to be extended (e.g. :class:`~pystac.Item`,
+    :class:`~pystac.Asset`).
+
+    To create a concrete instance of :class:`XarrayAssetsExtension`, use the
+    :meth:`XarrayAssetsExtension.ext` method. For example:
+
+    .. code-block:: python
+
+        >>> item: pystac.Item = ...
+        >>> xr_ext = XarrayAssetsExtension.ext(item)
+
+    """
+
+    @classmethod
+    def get_schema_uri(cls) -> str:
+        return SCHEMA_URI
+
+    @classmethod
+    def ext(cls, obj: T, add_if_missing: bool = False) -> XarrayAssetsExtension[T]:
+        """Extend the given STAC Object with properties from the
+        :stac-ext:`XarrayAssets Extension <xarray>`.
+
+        This extension can be applied to instances of :class:`~pystac.Collection`,
+        :class:`~pystac.Item` or :class:`~pystac.Asset`.
+
+        Raises:
+            pystac.ExtensionTypeError : If an invalid object type is passed.
+        """
+        if isinstance(obj, pystac.Collection):
+            cls.validate_has_extension(obj, add_if_missing)
+            return CollectionXarrayAssetsExtension(obj)
+        if isinstance(obj, pystac.Item):
+            cls.validate_has_extension(obj, add_if_missing)
+            return ItemXarrayAssetsExtension(obj)
+        if isinstance(obj, pystac.Asset):
+            cls.validate_owner_has_extension(obj, add_if_missing)
+            return AssetXarrayAssetsExtension(obj)
+        else:
+            raise pystac.ExtensionTypeError(cls._ext_error_message(obj))
+
+
+class CollectionXarrayAssetsExtension(XarrayAssetsExtension[pystac.Collection]):
+    """A concrete implementation of :class:`XarrayAssetsExtension` on a
+    :class:`~pystac.Collection` that extends the properties of the Item to include
+    properties defined in the :stac-ext:`XarrayAssets Extension <xarray>`.
+
+    This class should generally not be instantiated directly. Instead, call
+    :meth:`XarrayAssetsExtension.ext` on an :class:`~pystac.Collection` to extend it.
+    """
+
+    collection: pystac.Collection
+    properties: Dict[str, Any]
+
+    def __init__(self, collection: pystac.Collection):
+        self.collection = collection
+        self.properties = collection.extra_fields
+
+    def __repr__(self) -> str:
+        return "<CollectionXarrayAssetsExtension Item id={}>".format(self.collection.id)
+
+
+class ItemXarrayAssetsExtension(XarrayAssetsExtension[pystac.Item]):
+    """A concrete implementation of :class:`XarrayAssetsExtension` on an
+    :class:`~pystac.Item` that extends the properties of the Item to include properties
+    defined in the :stac-ext:`XarrayAssets Extension <xarray>`.
+
+    This class should generally not be instantiated directly. Instead, call
+    :meth:`XarrayAssetsExtension.ext` on an :class:`~pystac.Item` to extend it.
+    """
+
+    item: pystac.Item
+    properties: Dict[str, Any]
+
+    def __init__(self, item: pystac.Item):
+        self.item = item
+        self.properties = item.properties
+
+    def __repr__(self) -> str:
+        return "<ItemXarrayAssetsExtension Item id={}>".format(self.item.id)
+
+
+class AssetXarrayAssetsExtension(XarrayAssetsExtension[pystac.Asset]):
+    """A concrete implementation of :class:`XarrayAssetsExtension` on an
+    :class:`~pystac.Asset` that extends the Asset fields to include properties defined
+    in the :stac-ext:`XarrayAssets Extension <xarray>`.
+
+    This class should generally not be instantiated directly. Instead, call
+    :meth:`XarrayAssetsExtension.ext` on an :class:`~pystac.Asset` to extend it.
+    """
+
+    asset: pystac.Asset
+    properties: Dict[str, Any]
+    additional_read_properties: Optional[List[Dict[str, Any]]] = None
+
+    def __init__(self, asset: pystac.Asset):
+        self.asset = asset
+        self.properties = asset.extra_fields
+        if asset.owner and isinstance(asset.owner, pystac.Item):
+            self.additional_read_properties = [asset.owner.properties]
+
+    @property
+    def storage_options(self) -> Optional[Dict[str, Any]]:
+        """Additional keywords for accessing the dataset from remote storage"""
+        return self.properties.get(STORAGE_OPTIONS_PROP)
+
+    @storage_options.setter
+    def storage_options(self, v: Optional[Dict[str, Any]]) -> Any:
+        if v is None:
+            self.properties.pop(STORAGE_OPTIONS_PROP, None)
+        else:
+            self.properties[STORAGE_OPTIONS_PROP] = v
+
+    @property
+    def open_kwargs(self) -> Optional[Dict[str, Any]]:
+        """Additional keywords for opening the dataset"""
+        return self.properties.get(OPEN_KWARGS_PROP)
+
+    @open_kwargs.setter
+    def open_kwargs(self, v: Optional[Dict[str, Any]]) -> Any:
+        if v is None:
+            self.properties.pop(OPEN_KWARGS_PROP, None)
+        else:
+            self.properties[OPEN_KWARGS_PROP] = v
+
+    def __repr__(self) -> str:
+        return "<AssetXarrayAssetsExtension Asset href={}>".format(self.asset.href)
+
+
+class XarrayAssetsExtensionHooks(ExtensionHooks):
+    schema_uri: str = SCHEMA_URI
+    prev_extension_ids = {"xarray"}
+    stac_object_types = {pystac.STACObjectType.COLLECTION, pystac.STACObjectType.ITEM}
+
+
+XARRAY_ASSETS_EXTENSION_HOOKS: ExtensionHooks = XarrayAssetsExtensionHooks()
diff --git a/tests/data-files/xarray-assets/collection.json b/tests/data-files/xarray-assets/collection.json
@@ -0,0 +1,61 @@
+{
+    "stac_version": "1.0.0",
+    "stac_extensions": [
+      "https://stac-extensions.github.io/xarray-assets/v1.0.0/schema.json"
+    ],
+    "type": "Collection",
+    "id": "collection",
+    "title": "A title",
+    "description": "A description",
+    "license": "Apache-2.0",
+    "extent": {
+      "spatial": {
+        "bbox": [
+          [
+            172.9,
+            1.3,
+            173,
+            1.4
+          ]
+        ]
+      },
+      "temporal": {
+        "interval": [
+          [
+            "2015-06-23T00:00:00Z",
+            null
+          ]
+        ]
+      }
+    },
+    "assets": {
+      "example": {
+        "href": "abfs://cpdata/raw/terraclimate/4000m/raster.zarr",
+        "type": "application/vnd+zarr",
+        "xarray:storage_options": {
+          "account_name": "cpdataeuwest"
+        },
+        "xarray:open_kwargs": {
+          "consolidated": true
+        }
+      }
+    },
+    "summaries": {
+      "datetime": {
+        "minimum": "2015-06-23T00:00:00Z",
+        "maximum": "2019-07-10T13:44:56Z"
+      }
+    },
+    "links": [
+      {
+        "href": "./collection.json",
+        "rel": "root",
+        "title": "A title",
+        "type": "application/json"
+      },
+      {
+        "href": "./item.json",
+        "rel": "item"
+      }
+    ]
+  }
diff --git a/tests/data-files/xarray-assets/item.json b/tests/data-files/xarray-assets/item.json
@@ -0,0 +1,56 @@
+{
+    "stac_version": "1.0.0",
+    "stac_extensions": [
+      "https://stac-extensions.github.io/xarray-assets/v1.0.0/schema.json"
+    ],
+    "type": "Feature",
+    "id": "item",
+    "bbox": [
+      172.9,
+      1.3,
+      173,
+      1.4
+    ],
+    "geometry": {
+      "type": "Polygon",
+      "coordinates": [
+        [
+          [
+            172.9,
+            1.3
+          ],
+          [
+            173,
+            1.3
+          ],
+          [
+            173,
+            1.4
+          ],
+          [
+            172.9,
+            1.4
+          ],
+          [
+            172.9,
+            1.3
+          ]
+        ]
+      ]
+    },
+    "properties": {
+      "datetime": "2020-12-11T22:38:32Z"
+    },
+    "links": [],
+    "assets": {
+      "data": {
+        "href": "abfs://example.com/examples/file.zarr",
+        "xarray:storage_options": {
+          "account_name": "test-account"
+        },
+        "xarray:open_kwargs": {
+          "consolidated": true
+        }
+      }
+    }
+  }