[wip]

Author: jpolchlo

pystac/extensions/classification.py

@@ -23,6 +23,7 @@
 )
 from pystac.extensions.hooks import ExtensionHooks
 import pystac.extensions.item_assets as item_assets
+from pystac.extensions.raster import RasterExtension
 from pystac.serialization.identify import STACJSONDescription, STACVersionID
 from pystac.utils import get_required, map_opt
 
@@ -36,6 +37,7 @@
 CLASSES_PROP: str = PREFIX + "classes"
 RASTER_BANDS_PROP: str = "raster:bands"
 
+
 class Classification:
     properties: Dict[str, Any]
 
@@ -48,9 +50,7 @@ def apply(
         name: str,
         description: Optional[str] = None,
         title: Optional[str] = None,
-        color_hint: Optional[
-            str
-        ] = None,
+        color_hint: Optional[str] = None,
         nodata: Optional[bool] = None,
     ) -> None:
         self.value = value
@@ -62,7 +62,9 @@ def apply(
 
         color_hint_pattern = re.compile("^([0-9A-F]{6})$")
         match = color_hint_pattern.match(color_hint)
-        assert color_hint is None or match is not None and match.group() == color_hint, "Must format color hints as '^([0-9A-F]{6})$'"
+        assert (
+            color_hint is None or match is not None and match.group() == color_hint
+        ), "Must format color hints as '^([0-9A-F]{6})$'"
 
     @classmethod
     def create(
@@ -165,7 +167,7 @@ def apply(
         classes: List[Classification],
         roles: Optional[List[str]] = None,
         description: Optional[str] = None,
-        name: Optional[str] = None
+        name: Optional[str] = None,
     ) -> None:
         self.offset = offset
         self.length = length
@@ -177,11 +179,9 @@ def apply(
         assert offset >= 0, "Non-negative offsets only"
         assert length >= 1, "Positive field lengths only"
         assert len(classes) > 0, "Must specify at least one class"
-        assert roles is None or len(roles) > 0, "When set, roles must contain at least one item"
-
-        ## Ensure complete coverage of bit fields in classes
-        # class_coverage = set([c.value for c in classes])
-        # assert set(range(0, 2**length)) - class_coverage == set(), "Classes must cover the complete range of values"
+        assert (
+            roles is None or len(roles) > 0
+        ), "When set, roles must contain at least one item"
 
     @classmethod
     def create(
@@ -222,7 +222,10 @@ def length(self, v: int) -> None:
 
     @property
     def classes(self) -> List[Classification]:
-        return [Classification(d) for d in get_required(self.properties.get("classes"), self, "classes")]
+        return [
+            Classification(d)
+            for d in get_required(self.properties.get("classes"), self, "classes")
+        ]
 
     @classes.setter
     def classes(self, v: List[Classification]) -> None:
@@ -262,29 +265,24 @@ def name(self, v: Optional[str]) -> None:
             self.properties.pop("name", None)
 
     def __repr__(self) -> str:
-        return f"<Bitfield offset={self.offset} length={self.length} classes={self.classes}>"
+        return (
+            f"<Bitfield offset={self.offset} length={self.length} "
+            "classes={self.classes}>"
+        )
 
     def to_dict(self) -> Dict[str, Any]:
         return self.properties
 
 
 def __fetch_from_dict(d: Dict[str, Any]) -> Optional[Union[Classification, Bitfield]]:
-    assert not (CLASSES_PROP in d and BITFIELDS_PROP in d), f"Cannot define both {CLASSES_PROP} and {BITFIELDS:PROP} in the same entity"
+    assert not (
+        CLASSES_PROP in d and BITFIELDS_PROP in d
+    ), f"Cannot define both {CLASSES_PROP} and {BITFIELDS_PROP} in the same entity"
 
     if CLASSES_PROP in d:
-        return list(
-            map(
-                lambda item: Classification(item),
-                d[CLASSES_PROP]
-            )
-        )
+        return list(map(lambda item: Classification(item), d[CLASSES_PROP]))
     elif BITFIELDS_PROP in d:
-        return list(
-            map(
-                lambda item: Bitfield(item),
-                d[BITFIELDS_PROP]
-            )
-        )
+        return list(map(lambda item: Bitfield(item), d[BITFIELDS_PROP]))
     else:
         return None
 
@@ -299,7 +297,12 @@ def apply(
         classes: Optional[List[Classification]] = None,
         bitfields: Optional[List[Bitfield]] = None,
     ) -> None:
-        # assert classes is None and bitfields is not None or bitfields is None and classes is not None, "Must set exactly one of `classes` or `bitfields`"
+        assert (
+            classes is None
+            and bitfields is not None
+            or bitfields is None
+            and classes is not None
+        ), "Must set exactly one of `classes` or `bitfields`"
         self.classes = classes
         self.bitfields = bitfields
 
@@ -348,14 +351,17 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> ClassificationExtension[T]
         elif isinstance(obj, pystac.Asset):
             cls.validate_owner_has_extension(obj, add_if_missing)
             return cast(ClassificationExtension[T], AssetClassificationExtension(obj))
-        elif isinstance(obj, pystac.Collection):
-            cls.validate_has_extension(obj, add_if_missing)
-            return cast(ClassificationExtension[T], CollectionClassificationExtension(obj))
         elif isinstance(obj, item_assets.AssetDefinition):
-            return cast(ClassificationExtension[T], ItemAssetsClassificationExtension(obj))
+            cls.validate_has_extension(
+                cast(Union[pystac.Item, pystac.Collection], obj.owner), add_if_missing
+            )
+            return cast(
+                ClassificationExtension[T], ItemAssetsClassificationExtension(obj)
+            )
         else:
             raise pystac.ExtensionTypeError(
-                f"Classification extension does not apply to type '{type(obj).__name__}'"
+                "Classification extension does not apply to type "
+                f"'{type(obj).__name__}'"
             )
 
     @classmethod
@@ -383,11 +389,11 @@ def _get_classes(self) -> Optional[List[Classification]]:
 
         return None
 
-    def _get_bitfields(self) -> Optional[List[Bitfields]]:
+    def _get_bitfields(self) -> Optional[List[Bitfield]]:
         bitfields = self._get_property(BITFIELDS_PROP, List[Dict[str, Any]])
 
         if bitfields is not None:
-            return [Bitfields(b) for b in bitfields]
+            return [Bitfield(b) for b in bitfields]
 
         return None
 
@@ -403,13 +409,10 @@ class AssetClassificationExtension(ClassificationExtension[pystac.Asset]):
 
     def _get_classes(self) -> Optional[List[Classification]]:
         if CLASSES_PROP not in self.properties:
-            if (self.asset.owner and RasterExtension.has_extension(self.asset.owner)):
+            if self.asset.owner and RasterExtension.has_extension(self.asset.owner):
                 bnds = RasterExtension.ext(self.asset).bands
                 return list(
-                    map(
-                        lambda b: Classification(b.properties.get(CLASSES_PROP)),
-                        bnds
-                    )
+                    map(lambda b: Classification(b.properties.get(CLASSES_PROP)), bnds)
                 )
             else:
                 return None
@@ -418,13 +421,10 @@ def _get_classes(self) -> Optional[List[Classification]]:
 
     def _get_bitfields(self) -> Optional[List[Bitfield]]:
         if BITFIELDS_PROP not in self.properties:
-            if (self.asset.owner and RasterExtension.has_extension(self.asset.owner)):
+            if self.asset.owner and RasterExtension.has_extension(self.asset.owner):
                 bnds = RasterExtension.ext(self.asset).bands
                 return list(
-                    map(
-                        lambda b: Bitfield(b.properties.get(BITFIELDS_PROP)),
-                        bnds
-                    )
+                    map(lambda b: Bitfield(b.properties.get(BITFIELDS_PROP)), bnds)
                 )
             else:
                 return None
@@ -442,7 +442,9 @@ def __repr__(self) -> str:
         return f"<AssetClassificationExtension Asset href={self.asset_href}>"
 
 
-class ItemAssetsClassificationExtension(ClassificationExtension[item_assets.AssetDefinition]):
+class ItemAssetsClassificationExtension(
+    ClassificationExtension[item_assets.AssetDefinition]
+):
     properties: Dict[str, Any]
     asset_defn: item_assets.AssetDefinition
 
@@ -454,52 +456,30 @@ def __repr__(self) -> str:
         return f"<ItemAssetsClassificationExtension AssetDefinition={self.asset_defn}"
 
 
-class CollectionClassificationExtension(ClassificationExtension[pystac.Collection]):
-    """A concrete implementation of :class:`ClassificationExtension` on a
-    :class:`~pystac.Collection` that extends the properties of a Collection to
-    include properties defined in the :stac_ext:`Classification Extension
-    <classification>`.
-    """
-    collection: pystac.Collection
-    properties: Dict[str, Any]
+class SummariesClassificationExtension(SummariesExtension):
+    @property
+    def classes(self) -> Optional[List[Classification]]:
+        return map_opt(
+            lambda classes: [Classification(c) for c in classes],
+            self.summaries.get_list(CLASSES_PROP),
+        )
 
-    def __init__(self, collection: pystac.Collection):
-        self.collection = collection
-        self.properties = collection.assets
+    @classes.setter
+    def classes(self, v: Optional[List[Classification]]) -> None:
+        self._set_summary(CLASSES_PROP, map_opt(lambda x: [c.to_dict() for c in x], v))
 
-    def _get_classes(self) -> Optional[List[Optional[Classification]]]:
-        classes = __fetch_from_dict(self.properties)
-        if classes is None and ItemAssetsExtension.has_extension(self.collection):
-            classes = __fetch_from_dict(ItemAssetsExtension.ext(self.collection).item_assets)
-        return list(filter(lambda x: isinstance(x, Classification), classes))
+    @property
+    def bitfields(self) -> Optional[List[Bitfield]]:
+        return map_opt(
+            lambda bitfields: [Bitfield(b) for b in bitfields],
+            self.summaries.get_list(BITFIELDS_PROP),
+        )
 
-    def _get_bitfields(self) -> Optional[List[Bitfield]]:
-        pass
-
-# class SummariesClassificationExtension(SummariesExtension):
-#     @property
-#     def classes(self) -> Optional[List[Classification]]:
-#         return map_opt(
-#             lambda classes: [Classification(c) for c in classes],
-#             self.summaries.get_list(CLASSES_PROP),
-#         )
-
-#     @classes.setter
-#     def classes(self, v: Optional[List[Classification]]) -> None:
-#         self._set_summary(CLASSES_PROP, map_opt(lambda x: [c.to_dict() for c in x], v))
-
-#     @property
-#     def bitfields(self) -> Optional[List[Bitfield]]:
-#         return map_opt(
-#             lambda bitfields: [Bitfield(b) for b in bitfields],
-#             self.summaries.get_list(BITFIELDS_PROP),
-#         )
-
-#     @bitfields.setter
-#     def bitfields(self, v: Optional[List[Bitfield]]) -> None:
-#         self._set_summary(
-#             BITFIELDS_PROP, map_opt(lambda x: [b.to_dict() for b in x], v)
-#         )
+    @bitfields.setter
+    def bitfields(self, v: Optional[List[Bitfield]]) -> None:
+        self._set_summary(
+            BITFIELDS_PROP, map_opt(lambda x: [b.to_dict() for b in x], v)
+        )
 
 
 class ClassificationExtensionHooks(ExtensionHooks):