Fix item assets

Author: jsignell

pystac/extensions/base.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 import re
 import warnings
 from abc import ABC, abstractmethod
 from typing import (
+    TYPE_CHECKING,
     Any,
     Dict,
     Generic,
@@ -16,6 +19,9 @@
 
 import pystac
 
+if TYPE_CHECKING:
+    from pystac.extensions.item_assets import AssetDefinition
+
 VERSION_REGEX = re.compile("/v[0-9].[0-9].*/")
 
 
@@ -157,7 +163,39 @@ def has_extension(cls, obj: S) -> bool:
     @classmethod
     def validate_owner_has_extension(
         cls,
-        asset: pystac.Asset,
+        asset: Union[pystac.Asset, AssetDefinition],
+        add_if_missing: bool = False,
+    ) -> None:
+        """
+        DEPRECATED
+
+        .. deprecated:: 1.9
+            Use :meth:`ensure_owner_has_extension` instead.
+
+        Given an :class:`~pystac.Asset`, checks if the asset's owner has this
+        extension's schema URI in its :attr:`~pystac.STACObject.stac_extensions` list.
+        If ``add_if_missing`` is ``True``, the schema URI will be added to the owner.
+
+        Args:
+            asset : The asset whose owner should be validated.
+            add_if_missing : Whether to add the schema URI to the owner if it is
+                not already present. Defaults to False.
+
+        Raises:
+            STACError : If ``add_if_missing`` is ``True`` and ``asset.owner`` is
+                ``None``.
+        """
+        warnings.warn(
+            "ensure_owner_has_extension is deprecated and will be removed in v1.9. "
+            "Use ensure_owner_has_extension instead",
+            DeprecationWarning,
+        )
+        return cls.ensure_owner_has_extension(asset, add_if_missing)
+
+    @classmethod
+    def ensure_owner_has_extension(
+        cls,
+        asset: Union[pystac.Asset, AssetDefinition],
         add_if_missing: bool = False,
     ) -> None:
         """Given an :class:`~pystac.Asset`, checks if the asset's owner has this
@@ -176,8 +214,8 @@ def validate_owner_has_extension(
         if asset.owner is None:
             if add_if_missing:
                 raise pystac.STACError(
-                    "Attempted to use add_if_missing=True for an Asset with no owner. "
-                    "Use Asset.set_owner or set add_if_missing=False."
+                    "Attempted to use add_if_missing=True for an Asset or ItemAsset "
+                    "with no owner. Use .set_owner or set add_if_missing=False."
                 )
             else:
                 return

pystac/extensions/classification.py

@@ -19,7 +19,7 @@
 )
 
 import pystac
-import pystac.extensions.item_assets as item_assets
+from pystac.extensions import item_assets
 from pystac.extensions.base import (
     ExtensionManagementMixin,
     PropertiesExtension,
@@ -536,12 +536,10 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> ClassificationExtension[T]
             cls.ensure_has_extension(obj, add_if_missing)
             return cast(ClassificationExtension[T], ItemClassificationExtension(obj))
         elif isinstance(obj, pystac.Asset):
-            cls.validate_owner_has_extension(obj, add_if_missing)
+            cls.ensure_owner_has_extension(obj, add_if_missing)
             return cast(ClassificationExtension[T], AssetClassificationExtension(obj))
         elif isinstance(obj, item_assets.AssetDefinition):
-            cls.ensure_has_extension(
-                cast(Union[pystac.Item, pystac.Collection], obj.owner), add_if_missing
-            )
+            cls.ensure_owner_has_extension(obj, add_if_missing)
             return cast(
                 ClassificationExtension[T], ItemAssetsClassificationExtension(obj)
             )

pystac/extensions/datacube.py

@@ -6,11 +6,14 @@
 from typing import Any, Dict, Generic, List, Literal, Optional, TypeVar, Union, cast
 
 import pystac
+from pystac.extensions import item_assets
 from pystac.extensions.base import ExtensionManagementMixin, PropertiesExtension
 from pystac.extensions.hooks import ExtensionHooks
 from pystac.utils import StringEnum, get_required, map_opt
 
-T = TypeVar("T", pystac.Collection, pystac.Item, pystac.Asset)
+T = TypeVar(
+    "T", pystac.Collection, pystac.Item, pystac.Asset, item_assets.AssetDefinition
+)
 
 SCHEMA_URI = "https://stac-extensions.github.io/datacube/v2.0.0/schema.json"
 
@@ -545,8 +548,11 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> DatacubeExtension[T]:
             cls.ensure_has_extension(obj, add_if_missing)
             return cast(DatacubeExtension[T], ItemDatacubeExtension(obj))
         elif isinstance(obj, pystac.Asset):
-            cls.validate_owner_has_extension(obj, add_if_missing)
+            cls.ensure_owner_has_extension(obj, add_if_missing)
             return cast(DatacubeExtension[T], AssetDatacubeExtension(obj))
+        elif isinstance(obj, item_assets.AssetDefinition):
+            cls.ensure_owner_has_extension(obj, add_if_missing)
+            return cast(DatacubeExtension[T], ItemAssetsDatacubeExtension(obj))
         else:
             raise pystac.ExtensionTypeError(cls._ext_error_message(obj))
 
@@ -616,6 +622,15 @@ def __repr__(self) -> str:
         return "<AssetDatacubeExtension Item id={}>".format(self.asset_href)
 
 
+class ItemAssetsDatacubeExtension(DatacubeExtension[item_assets.AssetDefinition]):
+    properties: Dict[str, Any]
+    asset_defn: item_assets.AssetDefinition
+
+    def __init__(self, item_asset: item_assets.AssetDefinition):
+        self.asset_defn = item_asset
+        self.properties = item_asset.properties
+
+
 class DatacubeExtensionHooks(ExtensionHooks):
     schema_uri: str = SCHEMA_URI
     prev_extension_ids = {"datacube"}

pystac/extensions/eo.py

@@ -18,7 +18,7 @@
 )
 
 import pystac
-from pystac.extensions import projection, view
+from pystac.extensions import item_assets, projection, view
 from pystac.extensions.base import (
     ExtensionManagementMixin,
     PropertiesExtension,
@@ -29,7 +29,7 @@
 from pystac.summaries import RangeSummary
 from pystac.utils import get_required, map_opt
 
-T = TypeVar("T", pystac.Item, pystac.Asset)
+T = TypeVar("T", pystac.Item, pystac.Asset, item_assets.AssetDefinition)
 
 SCHEMA_URI: str = "https://stac-extensions.github.io/eo/v1.1.0/schema.json"
 SCHEMA_URIS: List[str] = [
@@ -411,8 +411,11 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> EOExtension[T]:
             cls.ensure_has_extension(obj, add_if_missing)
             return cast(EOExtension[T], ItemEOExtension(obj))
         elif isinstance(obj, pystac.Asset):
-            cls.validate_owner_has_extension(obj, add_if_missing)
+            cls.ensure_owner_has_extension(obj, add_if_missing)
             return cast(EOExtension[T], AssetEOExtension(obj))
+        elif isinstance(obj, item_assets.AssetDefinition):
+            cls.ensure_owner_has_extension(obj, add_if_missing)
+            return cast(EOExtension[T], ItemAssetsEOExtension(obj))
         else:
             raise pystac.ExtensionTypeError(cls._ext_error_message(obj))
 
@@ -537,6 +540,25 @@ def __repr__(self) -> str:
         return "<AssetEOExtension Asset href={}>".format(self.asset_href)
 
 
+class ItemAssetsEOExtension(EOExtension[item_assets.AssetDefinition]):
+    properties: Dict[str, Any]
+    asset_defn: item_assets.AssetDefinition
+
+    def _get_bands(self) -> Optional[List[Band]]:
+        if BANDS_PROP not in self.properties:
+            return None
+        return list(
+            map(
+                lambda band: Band(band),
+                cast(List[Dict[str, Any]], self.properties.get(BANDS_PROP)),
+            )
+        )
+
+    def __init__(self, item_asset: item_assets.AssetDefinition):
+        self.asset_defn = item_asset
+        self.properties = item_asset.properties
+
+
 class SummariesEOExtension(SummariesExtension):
     """A concrete implementation of :class:`~SummariesExtension` that extends
     the ``summaries`` field of a :class:`~pystac.Collection` to include properties

pystac/extensions/ext.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Dict, Generic, Literal, TypeVar, cast
+from typing import Any, Dict, Generic, Literal, TypeVar, cast
 
 import pystac
 from pystac.extensions.classification import ClassificationExtension
@@ -22,8 +22,7 @@
 from pystac.extensions.view import ViewExtension
 from pystac.extensions.xarray_assets import XarrayAssetsExtension
 
-if TYPE_CHECKING:
-    T = TypeVar("T", pystac.Asset, AssetDefinition)
+T = TypeVar("T", pystac.Asset, AssetDefinition)
 
 EXTENSION_NAMES = Literal[
     "classification",
@@ -192,9 +191,8 @@ def xarray(self) -> XarrayAssetsExtension[pystac.Item]:
         return XarrayAssetsExtension.ext(self.stac_object)
 
 
-@dataclass
-class AssetExt(Generic[T]):
-    stac_object: pystac.Asset
+class _AssetExt(Generic[T]):
+    stac_object: T
 
     def has(self, name: EXTENSION_NAMES) -> bool:
         if self.stac_object.owner is None:
@@ -237,10 +235,6 @@ def cube(self) -> DatacubeExtension[T]:
     def eo(self) -> EOExtension[T]:
         return EOExtension.ext(self.stac_object)
 
-    @property
-    def file(self) -> FileExtension:
-        return FileExtension.ext(self.stac_object)
-
     @property
     def pc(self) -> PointcloudExtension[T]:
         return PointcloudExtension.ext(self.stac_object)
@@ -269,14 +263,28 @@ def storage(self) -> StorageExtension[T]:
     def table(self) -> TableExtension[T]:
         return TableExtension.ext(self.stac_object)
 
-    @property
-    def timestamps(self) -> TimestampsExtension[T]:
-        return TimestampsExtension.ext(self.stac_object)
-
     @property
     def view(self) -> ViewExtension[T]:
         return ViewExtension.ext(self.stac_object)
 
+
+@dataclass
+class AssetExt(_AssetExt[pystac.Asset]):
+    stac_object: pystac.Asset
+
+    @property
+    def file(self) -> FileExtension:
+        return FileExtension.ext(self.stac_object)
+
     @property
-    def xarray(self) -> XarrayAssetsExtension[T]:
+    def timestamps(self) -> TimestampsExtension[pystac.Asset]:
+        return TimestampsExtension.ext(self.stac_object)
+
+    @property
+    def xarray(self) -> XarrayAssetsExtension[pystac.Asset]:
         return XarrayAssetsExtension.ext(self.stac_object)
+
+
+@dataclass
+class ItemAssetExt(_AssetExt[AssetDefinition]):
+    stac_object: AssetDefinition

pystac/extensions/file.py

@@ -226,7 +226,7 @@ def ext(cls, obj: pystac.Asset, add_if_missing: bool = False) -> FileExtension:
         This extension can be applied to instances of :class:`~pystac.Asset`.
         """
         if isinstance(obj, pystac.Asset):
-            cls.validate_owner_has_extension(obj, add_if_missing)
+            cls.ensure_owner_has_extension(obj, add_if_missing)
             return cls(obj)
         else:
             raise pystac.ExtensionTypeError(cls._ext_error_message(obj))

pystac/extensions/item_assets.py

@@ -12,7 +12,7 @@
 from pystac.utils import get_required
 
 if TYPE_CHECKING:
-    from pystac.extensions.ext import AssetExt
+    from pystac.extensions.ext import ItemAssetExt
 
 SCHEMA_URI = "https://stac-extensions.github.io/item-assets/v1.0.0/schema.json"
 
@@ -208,10 +208,10 @@ def create_asset(self, href: str) -> pystac.Asset:
         )
 
     @property
-    def ext(self) -> AssetExt:
-        from pystac.extensions.ext import AssetExt
+    def ext(self) -> ItemAssetExt:
+        from pystac.extensions.ext import ItemAssetExt
 
-        return AssetExt(stac_object=self)
+        return ItemAssetExt(stac_object=self)
 
 
 class ItemAssetsExtension(ExtensionManagementMixin[pystac.Collection]):

pystac/extensions/pointcloud.py

@@ -16,6 +16,7 @@
 )
 
 import pystac
+from pystac.extensions import item_assets
 from pystac.extensions.base import (
     ExtensionManagementMixin,
     PropertiesExtension,
@@ -25,7 +26,7 @@
 from pystac.summaries import RangeSummary
 from pystac.utils import StringEnum, get_required, map_opt
 
-T = TypeVar("T", pystac.Item, pystac.Asset)
+T = TypeVar("T", pystac.Item, pystac.Asset, item_assets.AssetDefinition)
 
 SCHEMA_URI: str = "https://stac-extensions.github.io/pointcloud/v1.0.0/schema.json"
 PREFIX: str = "pc:"
@@ -468,8 +469,11 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> PointcloudExtension[T]:
                 raise pystac.ExtensionTypeError(
                     "Pointcloud extension does not apply to Collection Assets."
                 )
-            cls.validate_owner_has_extension(obj, add_if_missing)
+            cls.ensure_owner_has_extension(obj, add_if_missing)
             return cast(PointcloudExtension[T], AssetPointcloudExtension(obj))
+        elif isinstance(obj, item_assets.AssetDefinition):
+            cls.ensure_owner_has_extension(obj, add_if_missing)
+            return cast(PointcloudExtension[T], ItemAssetsPointcloudExtension(obj))
         else:
             raise pystac.ExtensionTypeError(cls._ext_error_message(obj))
 
@@ -533,6 +537,15 @@ def __repr__(self) -> str:
         return f"<AssetPointcloudExtension Asset {self.repr_id}>"
 
 
+class ItemAssetsPointcloudExtension(PointcloudExtension[item_assets.AssetDefinition]):
+    properties: Dict[str, Any]
+    asset_defn: item_assets.AssetDefinition
+
+    def __init__(self, item_asset: item_assets.AssetDefinition):
+        self.asset_defn = item_asset
+        self.properties = item_asset.properties
+
+
 class SummariesPointcloudExtension(SummariesExtension):
     """A concrete implementation of :class:`~SummariesExtension` that extends
     the ``summaries`` field of a :class:`~pystac.Collection` to include properties

pystac/extensions/projection.py

@@ -18,14 +18,15 @@
 )
 
 import pystac
+from pystac.extensions import item_assets
 from pystac.extensions.base import (
     ExtensionManagementMixin,
     PropertiesExtension,
     SummariesExtension,
 )
 from pystac.extensions.hooks import ExtensionHooks
 
-T = TypeVar("T", pystac.Item, pystac.Asset)
+T = TypeVar("T", pystac.Item, pystac.Asset, item_assets.AssetDefinition)
 
 SCHEMA_URI: str = "https://stac-extensions.github.io/projection/v1.1.0/schema.json"
 SCHEMA_URIS: List[str] = [
@@ -301,8 +302,11 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> ProjectionExtension[T]:
             cls.ensure_has_extension(obj, add_if_missing)
             return cast(ProjectionExtension[T], ItemProjectionExtension(obj))
         elif isinstance(obj, pystac.Asset):
-            cls.validate_owner_has_extension(obj, add_if_missing)
+            cls.ensure_owner_has_extension(obj, add_if_missing)
             return cast(ProjectionExtension[T], AssetProjectionExtension(obj))
+        elif isinstance(obj, item_assets.AssetDefinition):
+            cls.ensure_owner_has_extension(obj, add_if_missing)
+            return cast(ProjectionExtension[T], ItemAssetsProjectionExtension(obj))
         else:
             raise pystac.ExtensionTypeError(cls._ext_error_message(obj))
 
@@ -367,6 +371,15 @@ def __repr__(self) -> str:
         return "<AssetProjectionExtension Asset href={}>".format(self.asset_href)
 
 
+class ItemAssetsProjectionExtension(ProjectionExtension[item_assets.AssetDefinition]):
+    properties: Dict[str, Any]
+    asset_defn: item_assets.AssetDefinition
+
+    def __init__(self, item_asset: item_assets.AssetDefinition):
+        self.asset_defn = item_asset
+        self.properties = item_asset.properties
+
+
 class SummariesProjectionExtension(SummariesExtension):
     """A concrete implementation of :class:`~SummariesExtension` that extends
     the ``summaries`` field of a :class:`~pystac.Collection` to include properties