feat: AssetDefinitions work with raster extension

Author: gadomski

CHANGELOG.md

@@ -10,6 +10,7 @@
 - MGRS Extension ([#1088](https://github.com/stac-utils/pystac/pull/1088))
 - All HTTP requests are logged when level is set to `logging.DEBUG` ([#1096](https://github.com/stac-utils/pystac/pull/1096))
 - `keep_parent` to Catalog `add_item` and `add_child` to avoid overriding existing parents ([#1117](https://github.com/stac-utils/pystac/pull/1117))
+- `owner` attribute to `AssetDefinition` in the item-assets extension ([#1110](https://github.com/stac-utils/pystac/pull/1110))
 
 ### Changed
 
@@ -21,6 +22,7 @@
 - `DefaultStacIO` now raises an error if it tries to write to a non-local url ([#1107](https://github.com/stac-utils/pystac/pull/1107))
 - Allow instantiation of pystac objects even with `"stac_extensions": null` ([#1109](https://github.com/stac-utils/pystac/pull/1109))
 - Make `Link.to_dict()` only contain strings ([#1114](https://github.com/stac-utils/pystac/pull/1114))
+- Updated raster extension to work with the item_assets extension's AssetDefinition objects ([#1110](https://github.com/stac-utils/pystac/pull/1110))
 
 ### Deprecated
 

pystac/extensions/item_assets.py

@@ -30,8 +30,13 @@ class AssetDefinition:
 
     properties: Dict[str, Any]
 
-    def __init__(self, properties: Dict[str, Any]) -> None:
+    owner: Optional[pystac.Collection]
+
+    def __init__(
+        self, properties: Dict[str, Any], owner: Optional[pystac.Collection] = None
+    ) -> None:
         self.properties = properties
+        self.owner = owner
 
     def __eq__(self, o: object) -> bool:
         if not isinstance(o, AssetDefinition):
@@ -107,6 +112,17 @@ def apply(
         self.description = description
         self.media_type = media_type
         self.roles = roles
+        self.owner = None
+
+    def set_owner(self, obj: pystac.Collection) -> None:
+        """Sets the owning item of this AssetDefinition.
+
+        The owning item will be used to resolve relative HREFs of this asset.
+
+        Args:
+            obj: The Collection that owns this asset.
+        """
+        self.owner = obj
 
     @property
     def title(self) -> Optional[str]:
@@ -202,7 +218,7 @@ def item_assets(self) -> Dict[str, AssetDefinition]:
         result: Dict[str, Any] = get_required(
             self.collection.extra_fields.get(ITEM_ASSETS_PROP), self, ITEM_ASSETS_PROP
         )
-        return {k: AssetDefinition(v) for k, v in result.items()}
+        return {k: AssetDefinition(v, self.collection) for k, v in result.items()}
 
     @item_assets.setter
     def item_assets(self, v: Dict[str, AssetDefinition]) -> None:

pystac/extensions/raster.py

@@ -2,7 +2,18 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, Iterable, List, Optional, Set, Union
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import pystac
 from pystac.extensions.base import (
@@ -11,8 +22,11 @@
     SummariesExtension,
 )
 from pystac.extensions.hooks import ExtensionHooks
+import pystac.extensions.item_assets as item_assets
 from pystac.utils import StringEnum, get_opt, get_required, map_opt
 
+T = TypeVar("T", pystac.Asset, item_assets.AssetDefinition)
+
 SCHEMA_URI = "https://stac-extensions.github.io/raster/v1.1.0/schema.json"
 SCHEMA_URIS = [
     "https://stac-extensions.github.io/raster/v1.0.0/schema.json",
@@ -644,38 +658,27 @@ def to_dict(self) -> Dict[str, Any]:
 
 
 class RasterExtension(
-    PropertiesExtension, ExtensionManagementMixin[Union[pystac.Item, pystac.Collection]]
+    Generic[T],
+    PropertiesExtension,
+    ExtensionManagementMixin[Union[pystac.Item, pystac.Collection]],
 ):
     """An abstract class that can be used to extend the properties of an
-    :class:`~pystac.Item` or :class:`~pystac.Asset` with properties from
+    :class:`~pystac.Item`, :class:`~pystac.Asset`, or
+    :class:`~pystac.extension.item_assets.AssetDefinition` with properties from
     the :stac-ext:`Raster Extension <raster>`. This class is generic over
     the type of STAC Object to be extended (e.g. :class:`~pystac.Item`,
     :class:`~pystac.Asset`).
 
     This class will generally not be used directly. Instead, use the concrete
     implementation associated with the STAC Object you want to extend (e.g.
-    :class:`~ItemRasterExtension` to extend an :class:`~pystac.Item`).
+    :class:`~ItemRasterExtension` to extend an :class:`~pystac.Item`).  You may
+    prefer to use the `ext` class method of this class to construct the correct
+    instance type for you.
     """
 
-    asset_href: str
-    """The ``href`` value of the :class:`~pystac.Asset` being extended."""
-
     properties: Dict[str, Any]
     """The :class:`~pystac.Asset` fields, including extension properties."""
 
-    additional_read_properties: Optional[Iterable[Dict[str, Any]]] = None
-    """If present, this will be a list containing 1 dictionary representing the
-    properties of the owning :class:`~pystac.Item`."""
-
-    def __init__(self, asset: pystac.Asset):
-        self.asset_href = asset.href
-        self.properties = asset.extra_fields
-        if asset.owner and isinstance(asset.owner, pystac.Item):
-            self.additional_read_properties = [asset.owner.properties]
-
-    def __repr__(self) -> str:
-        return "<AssetRasterExtension Asset href={}>".format(self.asset_href)
-
     def apply(self, bands: List[RasterBand]) -> None:
         """Applies raster extension properties to the extended :class:`pystac.Item` or
         :class:`pystac.Asset`.
@@ -715,7 +718,7 @@ def get_schema_uris(cls) -> List[str]:
         return SCHEMA_URIS
 
     @classmethod
-    def ext(cls, obj: pystac.Asset, add_if_missing: bool = False) -> RasterExtension:
+    def ext(cls, obj: T, add_if_missing: bool = False) -> RasterExtension[T]:
         """Extends the given STAC Object with properties from the :stac-ext:`Raster
         Extension <raster>`.
 
@@ -727,7 +730,12 @@ def ext(cls, obj: pystac.Asset, add_if_missing: bool = False) -> RasterExtension
         """
         if isinstance(obj, pystac.Asset):
             cls.validate_owner_has_extension(obj, add_if_missing)
-            return cls(obj)
+            return cast(RasterExtension[T], AssetRasterExtension(obj))
+        elif isinstance(obj, item_assets.AssetDefinition):
+            cls.validate_has_extension(
+                cast(Union[pystac.Item, pystac.Collection], obj.owner), add_if_missing
+            )
+            return cast(RasterExtension[T], ItemAssetsRasterExtension(obj))
         else:
             raise pystac.ExtensionTypeError(
                 f"Raster extension does not apply to type '{type(obj).__name__}'"
@@ -741,6 +749,46 @@ def summaries(
         return SummariesRasterExtension(obj)
 
 
+class AssetRasterExtension(RasterExtension[pystac.Asset]):
+    asset_href: str
+    """The ``href`` value of the :class:`~pystac.Asset` being extended."""
+
+    properties: Dict[str, Any]
+    """The :class:`~pystac.Asset` fields, including extension properties."""
+
+    additional_read_properties: Optional[Iterable[Dict[str, Any]]] = None
+    """If present, this will be a list containing 1 dictionary representing the
+    properties of the owning :class:`~pystac.Item`."""
+
+    def __init__(self, asset: pystac.Asset):
+        self.asset_href = asset.href
+        self.properties = asset.extra_fields
+        if asset.owner and isinstance(asset.owner, pystac.Item):
+            self.additional_read_properties = [asset.owner.properties]
+
+    def __repr__(self) -> str:
+        return "<AssetRasterExtension Asset href={}>".format(self.asset_href)
+
+
+class ItemAssetsRasterExtension(RasterExtension[item_assets.AssetDefinition]):
+    asset_definition: item_assets.AssetDefinition
+    """A reference to the :class:`~pystac.extensions.item_assets.AssetDefinition`
+    being extended."""
+
+    properties: Dict[str, Any]
+    """The :class:`~pystac.extensions.item_assets.AssetDefinition` fields, including
+    extension properties."""
+
+    def __init__(self, item_asset: item_assets.AssetDefinition):
+        self.properties = item_asset.properties
+        self.asset_definition = item_asset
+
+    def __repr__(self) -> str:
+        return "<ItemAssetsRasterExtension AssetDefinition={}>".format(
+            self.asset_definition
+        )
+
+
 class SummariesRasterExtension(SummariesExtension):
     """A concrete implementation of :class:`~SummariesExtension` that extends
     the ``summaries`` field of a :class:`~pystac.Collection` to include properties