Add collection to ExtensionManagementMixin type params where appropriate (#547)

Fix typing of extensions summaries and validate extension URI presence

Author: jisantuc

CHANGELOG.md

@@ -14,6 +14,10 @@
 
 - The `from_dict` method on STACObjects will set the object's root link when a `root` parameter is present. An ItemCollection `from_dict` with a root parameter will set the root on each of it's Items. ([#549](https://github.com/stac-utils/pystac/pull/549))
 
+### Fixed
+
+- Added `Collections` as a type that can be extended for extensions whose fields can appear in collection summaries ([#547](https://github.com/stac-utils/pystac/pull/547))
+
 ### Deprecated
 
 ## [v1.0.0-rc.3]

pystac/extensions/eo.py

@@ -12,6 +12,7 @@
     Optional,
     Tuple,
     TypeVar,
+    Union,
     cast,
 )
 
@@ -274,7 +275,9 @@ def band_description(common_name: str) -> Optional[str]:
 
 
 class EOExtension(
-    Generic[T], PropertiesExtension, ExtensionManagementMixin[pystac.Item]
+    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 the
@@ -296,7 +299,7 @@ def apply(
         self, bands: Optional[List[Band]] = None, cloud_cover: Optional[float] = None
     ) -> None:
         """Applies label extension properties to the extended :class:`~pystac.Item` or
-        :class:`~pystac.Collection`.
+        :class:`~pystac.Asset`.
 
         Args:
             bands : A list of available bands where each item is a :class:`~Band`
@@ -373,9 +376,15 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> "EOExtension[T]":
                 f"EO extension does not apply to type '{type(obj).__name__}'"
             )
 
-    @staticmethod
-    def summaries(obj: pystac.Collection) -> "SummariesEOExtension":
+    @classmethod
+    def summaries(
+        cls, obj: pystac.Collection, add_if_missing: bool = False
+    ) -> "SummariesEOExtension":
         """Returns the extended summaries object for the given collection."""
+        if not add_if_missing:
+            cls.validate_has_extension(obj)
+        else:
+            cls.add_to(obj)
         return SummariesEOExtension(obj)
 
 

pystac/extensions/label.py

@@ -425,7 +425,7 @@ def __eq__(self, o: object) -> bool:
         return self.to_dict() == o
 
 
-class LabelExtension(ExtensionManagementMixin[pystac.Item]):
+class LabelExtension(ExtensionManagementMixin[Union[pystac.Item, pystac.Collection]]):
     """A class that can be used to extend the properties of an
     :class:`~pystac.Item` with properties from the :stac-ext:`Label Extension <label>`.
 
@@ -700,9 +700,15 @@ def ext(cls, obj: pystac.Item, add_if_missing: bool = False) -> "LabelExtension"
                 f"Label extension does not apply to type '{type(obj).__name__}'"
             )
 
-    @staticmethod
-    def summaries(obj: pystac.Collection) -> "SummariesLabelExtension":
+    @classmethod
+    def summaries(
+        cls, obj: pystac.Collection, add_if_missing: bool = False
+    ) -> "SummariesLabelExtension":
         """Returns the extended summaries object for the given collection."""
+        if not add_if_missing:
+            cls.validate_has_extension(obj)
+        else:
+            cls.add_to(obj)
         return SummariesLabelExtension(obj)
 
 

pystac/extensions/projection.py

@@ -4,7 +4,7 @@
 """
 
 import json
-from typing import Any, Dict, Generic, Iterable, List, Optional, TypeVar, cast
+from typing import Any, Dict, Generic, Iterable, List, Optional, TypeVar, Union, cast
 
 import pystac
 from pystac.extensions.hooks import ExtensionHooks
@@ -31,7 +31,9 @@
 
 
 class ProjectionExtension(
-    Generic[T], PropertiesExtension, ExtensionManagementMixin[pystac.Item]
+    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` with properties from the :stac-ext:`Projection
@@ -284,9 +286,15 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> "ProjectionExtension[T]":
                 f"Projection extension does not apply to type '{type(obj).__name__}'"
             )
 
-    @staticmethod
-    def summaries(obj: pystac.Collection) -> "SummariesProjectionExtension":
+    @classmethod
+    def summaries(
+        cls, obj: pystac.Collection, add_if_missing: bool = False
+    ) -> "SummariesProjectionExtension":
         """Returns the extended summaries object for the given collection."""
+        if not add_if_missing:
+            cls.validate_has_extension(obj)
+        else:
+            cls.add_to(obj)
         return SummariesProjectionExtension(obj)
 
 

pystac/extensions/raster.py

@@ -4,7 +4,7 @@
 """
 
 import enum
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 import pystac
 from pystac.extensions.base import (
@@ -625,7 +625,9 @@ def to_dict(self) -> Dict[str, Any]:
         return self.properties
 
 
-class RasterExtension(PropertiesExtension, ExtensionManagementMixin[pystac.Item]):
+class RasterExtension(
+    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
     the :stac-ext:`Raster Extension <raster>`. This class is generic over
@@ -711,8 +713,14 @@ def ext(cls, obj: pystac.Asset, add_if_missing: bool = False) -> "RasterExtensio
                 f"Raster extension does not apply to type '{type(obj).__name__}'"
             )
 
-    @staticmethod
-    def summaries(obj: pystac.Collection) -> "SummariesRasterExtension":
+    @classmethod
+    def summaries(
+        cls, obj: pystac.Collection, add_if_missing: bool = False
+    ) -> "SummariesRasterExtension":
+        if not add_if_missing:
+            cls.validate_has_extension(obj)
+        else:
+            cls.add_to(obj)
         return SummariesRasterExtension(obj)
 
 

pystac/extensions/sat.py

@@ -6,7 +6,7 @@
 import enum
 from datetime import datetime as Datetime
 from pystac.summaries import RangeSummary
-from typing import Dict, Any, List, Iterable, Generic, Optional, TypeVar, cast
+from typing import Dict, Any, List, Iterable, Generic, Optional, TypeVar, Union, cast
 
 import pystac
 from pystac.extensions.base import (
@@ -38,7 +38,9 @@ class OrbitState(str, enum.Enum):
 
 
 class SatExtension(
-    Generic[T], PropertiesExtension, ExtensionManagementMixin[pystac.Item]
+    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 the
@@ -161,9 +163,15 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> "SatExtension[T]":
                 f"Satellite extension does not apply to type '{type(obj).__name__}'"
             )
 
-    @staticmethod
-    def summaries(obj: pystac.Collection) -> "SummariesSatExtension":
+    @classmethod
+    def summaries(
+        cls, obj: pystac.Collection, add_if_missing: bool = False
+    ) -> "SummariesSatExtension":
         """Returns the extended summaries object for the given collection."""
+        if not add_if_missing:
+            cls.validate_has_extension(obj)
+        else:
+            cls.add_to(obj)
         return SummariesSatExtension(obj)
 
 

pystac/extensions/scientific.py

@@ -248,9 +248,15 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> "ScientificExtension[T]":
                 f"Scientific extension does not apply to type '{type(obj).__name__}'"
             )
 
-    @staticmethod
-    def summaries(obj: pystac.Collection) -> "SummariesScientificExtension":
+    @classmethod
+    def summaries(
+        cls, obj: pystac.Collection, add_if_missing: bool = False
+    ) -> "SummariesScientificExtension":
         """Returns the extended summaries object for the given collection."""
+        if not add_if_missing:
+            cls.validate_has_extension(obj)
+        else:
+            cls.add_to(obj)
         return SummariesScientificExtension(obj)
 
 

pystac/extensions/timestamps.py

@@ -5,7 +5,7 @@
 
 from datetime import datetime as datetime
 from pystac.summaries import RangeSummary
-from typing import Dict, Any, Iterable, Generic, Optional, TypeVar, cast
+from typing import Dict, Any, Iterable, Generic, Optional, TypeVar, Union, cast
 
 import pystac
 from pystac.extensions.base import (
@@ -26,7 +26,9 @@
 
 
 class TimestampsExtension(
-    Generic[T], PropertiesExtension, ExtensionManagementMixin[pystac.Item]
+    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 the
@@ -141,9 +143,15 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> "TimestampsExtension[T]":
                 f"Timestamps extension does not apply to type '{type(obj).__name__}'"
             )
 
-    @staticmethod
-    def summaries(obj: pystac.Collection) -> "SummariesTimestampsExtension":
+    @classmethod
+    def summaries(
+        cls, obj: pystac.Collection, add_if_missing: bool = False
+    ) -> "SummariesTimestampsExtension":
         """Returns the extended summaries object for the given collection."""
+        if not add_if_missing:
+            cls.validate_has_extension(obj)
+        else:
+            cls.add_to(obj)
         return SummariesTimestampsExtension(obj)
 
 

pystac/extensions/view.py

@@ -3,7 +3,7 @@
 https://github.com/stac-extensions/view
 """
 
-from typing import Any, Dict, Generic, Iterable, Optional, TypeVar, cast
+from typing import Any, Dict, Generic, Iterable, Optional, TypeVar, Union, cast
 
 import pystac
 from pystac.summaries import RangeSummary
@@ -27,7 +27,9 @@
 
 
 class ViewExtension(
-    Generic[T], PropertiesExtension, ExtensionManagementMixin[pystac.Item]
+    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` with properties from the :stac-ext:`View Geometry
@@ -170,9 +172,15 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> "ViewExtension[T]":
                 f"View extension does not apply to type '{type(obj).__name__}'"
             )
 
-    @staticmethod
-    def summaries(obj: pystac.Collection) -> "SummariesViewExtension":
+    @classmethod
+    def summaries(
+        cls, obj: pystac.Collection, add_if_missing: bool = False
+    ) -> "SummariesViewExtension":
         """Returns the extended summaries object for the given collection."""
+        if not add_if_missing:
+            cls.validate_has_extension(obj)
+        else:
+            cls.add_to(obj)
         return SummariesViewExtension(obj)
 
 

tests/extensions/test_eo.py

@@ -214,6 +214,23 @@ def test_summaries(self) -> None:
         self.assertEqual(len(col_dict["summaries"]["eo:bands"]), 1)
         self.assertEqual(col_dict["summaries"]["eo:cloud_cover"]["minimum"], 1.0)
 
+    def test_summaries_adds_uri(self) -> None:
+        col = pystac.Collection.from_file(self.EO_COLLECTION_URI)
+        col.stac_extensions = []
+        self.assertRaisesRegex(
+            pystac.ExtensionNotImplemented,
+            r"Could not find extension schema URI.*",
+            EOExtension.summaries,
+            col,
+            False,
+        )
+        _ = EOExtension.summaries(col, add_if_missing=True)
+
+        self.assertIn(EOExtension.get_schema_uri(), col.stac_extensions)
+
+        EOExtension.remove_from(col)
+        self.assertNotIn(EOExtension.get_schema_uri(), col.stac_extensions)
+
     def test_read_pre_09_fields_into_common_metadata(self) -> None:
         eo_item = pystac.Item.from_file(
             TestCases.get_path(