Source code for pystac.extensions.xarray_assets
"""Implements the :stac-ext:`Xarray Assets Extension <xarray>`."""
from __future__ import annotations
from typing import Any, Generic, Literal, 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"
[docs]
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)
"""
name: Literal["xarray"] = "xarray"
[docs]
@classmethod
def get_schema_uri(cls) -> str:
return SCHEMA_URI
[docs]
@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.ensure_has_extension(obj, add_if_missing)
return CollectionXarrayAssetsExtension(obj)
if isinstance(obj, pystac.Item):
cls.ensure_has_extension(obj, add_if_missing)
return ItemXarrayAssetsExtension(obj)
if isinstance(obj, pystac.Asset):
cls.ensure_owner_has_extension(obj, add_if_missing)
return AssetXarrayAssetsExtension(obj)
else:
raise pystac.ExtensionTypeError(cls._ext_error_message(obj))
[docs]
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 f"<CollectionXarrayAssetsExtension Item id={self.collection.id}>"
[docs]
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 f"<ItemXarrayAssetsExtension Item id={self.item.id}>"
[docs]
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: list[dict[str, Any]] | None = 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) -> dict[str, Any] | None:
"""Additional keywords for accessing the dataset from remote storage"""
return self.properties.get(STORAGE_OPTIONS_PROP)
@storage_options.setter
def storage_options(self, v: dict[str, Any] | None) -> Any:
if v is None:
self.properties.pop(STORAGE_OPTIONS_PROP, None)
else:
self.properties[STORAGE_OPTIONS_PROP] = v
@property
def open_kwargs(self) -> dict[str, Any] | None:
"""Additional keywords for opening the dataset"""
return self.properties.get(OPEN_KWARGS_PROP)
@open_kwargs.setter
def open_kwargs(self, v: dict[str, Any] | None) -> Any:
if v is None:
self.properties.pop(OPEN_KWARGS_PROP, None)
else:
self.properties[OPEN_KWARGS_PROP] = v
def __repr__(self) -> str:
return f"<AssetXarrayAssetsExtension Asset href={self.asset.href}>"
[docs]
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()