Simpler extension interface

docs/concepts.rst

@@ -475,56 +475,56 @@ Extension (e.g. Electro-Optical, Projection, etc.) and STAC Object
 (:class:`~pystac.Collection`, :class:`pystac.Item`, or :class:`pystac.Asset`). All
 classes that extend these objects inherit from
 :class:`pystac.extensions.base.PropertiesExtension`, and you can use the
-``ext`` method on these classes to extend an object.
+``ext`` accessor on the object to access the extension fields.
 
 For instance, if you have an item that implements the :stac-ext:`Electro-Optical
-Extension <eo>`, you can access the properties associated with that extension using
-:meth:`EOExtension.ext <pystac.extensions.eo.EOExtension.ext>`:
+Extension <eo>`, you can access the fields associated with that extension using
+:meth:`Item.ext <pystac.Item.ext>`:
 
 .. code-block:: python
 
    import pystac
-   from pystac.extensions.eo import EOExtension
 
    item = pystac.Item.from_file("tests/data-files/eo/eo-landsat-example.json")
 
-   # Check that the Item implements the EO Extension
-   if EOExtension.has_extension(item):
-      eo_ext = EOExtension.ext(item)
+   # As long as the Item implements the EO Extension you can access all the
+   # EO properties directly
+   bands = item.ext.eo.bands
+   cloud_cover = item.ext.eo.cloud_cover
+   ...
 
-      bands = eo_ext.bands
-      cloud_cover = eo_ext.cloud_cover
-      snow_cover = eo_ext.snow_cover
-      ...
-
-.. note:: The ``ext`` method will raise an :exc:`~pystac.ExtensionNotImplemented`
+.. note:: ``ext`` will raise an :exc:`~pystac.ExtensionNotImplemented`
    exception if the object does not implement that extension (e.g. if the extension
-   URI is not in that object's :attr:`~pystac.STACObject.stac_extensions` list). In
-   the example above, we check that the Item implements the EO Extension before calling
-   :meth:`EOExtension.ext <pystac.extensions.eo.EOExtension.ext>` to handle this. See
+   URI is not in that object's :attr:`~pystac.STACObject.stac_extensions` list). See
    the `Adding an Extension`_ section below for details on adding an extension to an
    object.
 
+If you don't want to raise an error you can use :meth:`~pystac.Item.ext.has`
+to first check if the extension is implemented on your pystac object:
+
+.. code-block:: python
+
+   if item.ext.has("eo"):
+      bands = item.ext.eo.bands
+
 See the documentation for each extension implementation for details on the supported
 properties and other functionality.
 
-Instances of :class:`~pystac.extensions.base.PropertiesExtension` have a
-:attr:`~pystac.extensions.base.PropertiesExtension.properties` attribute that gives
-access to the properties of the extended object. *This attribute is a reference to the
-properties of the* :class:`~pystac.Item` *or* :class:`~pystac.Asset` *being extended and
-can therefore mutate those properties.* For instance:
+Extensions have access to the properties of the object. *This attribute is a reference
+to the properties of the* :class:`~pystac.Collection`, :class:`~pystac.Item` *or*
+:class:`~pystac.Asset` *being extended and can therefore mutate those properties.*
+For instance:
 
 .. code-block:: python
 
    item = pystac.Item.from_file("tests/data-files/eo/eo-landsat-example.json")
    print(item.properties["eo:cloud_cover"])
    # 78
 
-   eo_ext = EOExtension.ext(item)
-   print(eo_ext.cloud_cover)
+   print(item.ext.eo.cloud_cover)
    # 78
 
-   eo_ext.cloud_cover = 45
+   item.ext.eo.cloud_cover = 45
    print(item.properties["eo:cloud_cover"])
    # 45
 
@@ -545,24 +545,17 @@ have a default value of ``None``:
 .. code-block:: python
 
    # Can also omit cloud_cover entirely...
-   eo_ext.apply(0.5, bands, cloud_cover=None)
-
-
-If you attempt to extend an object that is not supported by an extension, PySTAC will
-throw a :class:`pystac.ExtensionTypeError`.
+   item.ext.eo.apply(0.5, bands, cloud_cover=None)
 
 
 Adding an Extension
 -------------------
 
 You can add an extension to a STAC object that does not already implement that extension
-using the :meth:`ExtensionManagementMixin.add_to
-<pystac.extensions.base.ExtensionManagementMixin.add_to>` method. Any concrete
-extension implementations that extend existing STAC objects should inherit from the
-:class:`~pystac.extensions.base.ExtensionManagementMixin` class, and will therefore have
-this method available. The
-:meth:`~pystac.extensions.base.ExtensionManagementMixin.add_to` adds the correct schema
-URI to the :attr:`~pystac.STACObject.stac_extensions` list for the object being
+using the :meth:`~pystac.Item.ext.add` method. Any concrete
+extension implementations that extend existing STAC objects should  have
+this method available. The :meth:`~pystac.Item.ext.add` method adds the correct schema
+URI to the :attr:`~pystac.Item.stac_extensions` list for the object being
 extended.
 
 .. code-block:: python
@@ -573,7 +566,7 @@ extended.
    # []
 
    # Add the Electro-Optical extension
-   EOExtension.add_to(item)
+   item.ext.add("eo")
    print(item.stac_extensions)
    # ['https://stac-extensions.github.io/eo/v1.1.0/schema.json']
 
@@ -617,7 +610,7 @@ Item Asset properties
 =====================
 
 Properties that apply to Items can be found in two places: the Item's properties or in
-any of an Item's Assets. If the property is on an Asset, it applies only that specific
+any of an Item's Assets. If the property is on an Asset, it applies only to that specific
 asset. For example, gsd defined for an Item represents the best Ground Sample Distance
 (resolution) for the data within the Item. However, some assets may be lower resolution
 and thus have a higher gsd. In that case, the `gsd` can be found on the Asset.

pystac/asset.py

@@ -12,6 +12,7 @@
 if TYPE_CHECKING:
     from pystac.collection import Collection
     from pystac.common_metadata import CommonMetadata
+    from pystac.extensions.ext import AssetExt
     from pystac.item import Item
 
 A = TypeVar("A", bound="Asset")
@@ -261,6 +262,12 @@ def delete(self) -> None:
         href = _absolute_href(self.href, self.owner, "delete")
         os.remove(href)
 
+    @property
+    def ext(self) -> AssetExt:
+        from pystac.extensions.ext import AssetExt
+
+        return AssetExt(stac_object=self)
+
 
 def _absolute_href(
     href: str, owner: Optional[Union[Item, Collection]], action: str = "access"

pystac/collection.py

@@ -42,6 +42,7 @@
 )
 
 if TYPE_CHECKING:
+    from pystac.extensions.ext import CollectionExt
     from pystac.item import Item
 
 C = TypeVar("C", bound="Collection")
@@ -830,3 +831,9 @@ def full_copy(
     @classmethod
     def matches_object_type(cls, d: Dict[str, Any]) -> bool:
         return identify_stac_object_type(d) == STACObjectType.COLLECTION
+
+    @property
+    def ext(self) -> CollectionExt:
+        from pystac.extensions.ext import CollectionExt
+
+        return CollectionExt(stac_object=self)

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
@@ -223,8 +261,16 @@ def ensure_has_extension(cls, obj: S, add_if_missing: bool = False) -> None:
             cls.add_to(obj)
 
         if not cls.has_extension(obj):
+            name = getattr(cls, "name", cls.__name__)
+            suggestion = (
+                f"``obj.ext.add('{name}')"
+                if hasattr(cls, "name")
+                else f"``{name}.add_to(obj)``"
+            )
+
             raise pystac.ExtensionNotImplemented(
-                f"Could not find extension schema URI {cls.get_schema_uri()} in object."
+                f"Extension '{name}' is not implemented on object."
+                f"STAC producers can add the extension using {suggestion}"
             )
 
     @classmethod

pystac/extensions/classification.py

@@ -10,6 +10,7 @@
     Generic,
     Iterable,
     List,
+    Literal,
     Optional,
     Pattern,
     TypeVar,
@@ -18,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,
@@ -426,6 +427,7 @@ class ClassificationExtension(
     method can be used to construct the proper class for you.
     """
 
+    name: Literal["classification"] = "classification"
     properties: Dict[str, Any]
     """The :class:`~pystac.Asset` fields, including extension properties."""
 
@@ -534,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

@@ -3,14 +3,17 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import Any, Dict, Generic, List, Optional, TypeVar, Union, cast
+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"
 
@@ -469,6 +472,8 @@ class DatacubeExtension(
        >>> dc_ext = DatacubeExtension.ext(item)
     """
 
+    name: Literal["cube"] = "cube"
+
     def apply(
         self,
         dimensions: Dict[str, Dimension],
@@ -543,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))
 
@@ -614,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"}