[wip]

Author: jpolchlo

pystac/extensions/classification.py

@@ -22,10 +22,11 @@
     SummariesExtension,
 )
 from pystac.extensions.hooks import ExtensionHooks
+import pystac.extensions.item_assets as item_assets
 from pystac.serialization.identify import STACJSONDescription, STACVersionID
 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/classification/v1.0.0/schema.json"
 PREFIX: str = "classification:"
@@ -350,6 +351,8 @@ def ext(cls, obj: T, add_if_missing: bool = False) -> ClassificationExtension[T]
         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))
         else:
             raise pystac.ExtensionTypeError(
                 f"Classification extension does not apply to type '{type(obj).__name__}'"
@@ -439,6 +442,18 @@ def __repr__(self) -> str:
         return f"<AssetClassificationExtension Asset href={self.asset_href}>"
 
 
+class ItemAssetsClassificationExtension(ClassificationExtension[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
+
+    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
@@ -461,30 +476,30 @@ def _get_classes(self) -> Optional[List[Optional[Classification]]]:
     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)
-        )
+# 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)
+#         )
 
 
 class ClassificationExtensionHooks(ExtensionHooks):