Add filter to Item and Collection get_assets method (#936)

* feat: filter Item and Collection get_assets

* chore: update CHANGELOG

* refactor: reduce iterations over asset dictionary

* refactor: simplify role checking logic

* chore: update CHANGELOG

Co-authored-by: Pete Gadomski <pete.gadomski@gmail.com>

* fix: correct typing

Co-authored-by: Pete Gadomski <pete.gadomski@gmail.com>

Author: pjhartzell

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Added
 
+- Add `media_type` and `role` filtering to Item and Collection `get_assets()` method ([#936](https://github.com/stac-utils/pystac/pull/936))
+- `Asset.has_role` ([#936](https://github.com/stac-utils/pystac/pull/936))
 - Enum MediaType entry for flatgeobuf ([discussion](https://github.com/flatgeobuf/flatgeobuf/discussions/112#discussioncomment-4606721)) ([#938](https://github.com/stac-utils/pystac/pull/938))
 - Update Grid Extension support to v1.1.0 and fix issue with grid:code prefix validation ([#925](https://github.com/stac-utils/pystac/pull/925))
 - Adds custom `header` support to `DefaultStacIO` ([#889](https://github.com/stac-utils/pystac/pull/889))

pystac/asset.py

@@ -152,6 +152,20 @@ def clone(self) -> "Asset":
             extra_fields=deepcopy(self.extra_fields),
         )
 
+    def has_role(self, role: str) -> bool:
+        """Check if a role exists in the Asset role list.
+
+        Args:
+            role: Role to check for existence.
+
+        Returns:
+            bool: True if role exists, else False.
+        """
+        if self.roles is None:
+            return False
+        else:
+            return role in self.roles
+
     @property
     def common_metadata(self) -> "CommonMetadata_Type":
         """Access the asset's common metadata fields as a

pystac/collection.py

@@ -675,13 +675,32 @@ def from_dict(
 
         return collection
 
-    def get_assets(self) -> Dict[str, Asset]:
-        """Get this item's assets.
+    def get_assets(
+        self,
+        media_type: Optional[Union[str, pystac.MediaType]] = None,
+        role: Optional[str] = None,
+    ) -> Dict[str, Asset]:
+        """Get this collection's assets.
+
+        Args:
+            media_type: If set, filter the assets such that only those with a
+                matching ``media_type`` are returned.
+            role: If set, filter the assets such that only those with a matching
+                ``role`` are returned.
 
         Returns:
-            Dict[str, Asset]: A copy of the dictionary of this item's assets.
+            Dict[str, Asset]: A dictionary of assets that match ``media_type``
+                and/or ``role`` if set or else all of this collection's assets.
         """
-        return dict(self.assets.items())
+        if media_type is None and role is None:
+            return dict(self.assets.items())
+        assets = dict()
+        for key, asset in self.assets.items():
+            if (media_type is None or asset.media_type == media_type) and (
+                role is None or asset.has_role(role)
+            ):
+                assets[key] = asset
+        return assets
 
     def add_asset(self, key: str, asset: Asset) -> None:
         """Adds an Asset to this item.

pystac/item.py

@@ -221,13 +221,32 @@ def set_datetime(self, datetime: Datetime, asset: Optional[Asset] = None) -> Non
         else:
             asset.extra_fields["datetime"] = datetime_to_str(datetime)
 
-    def get_assets(self) -> Dict[str, Asset]:
+    def get_assets(
+        self,
+        media_type: Optional[Union[str, pystac.MediaType]] = None,
+        role: Optional[str] = None,
+    ) -> Dict[str, Asset]:
         """Get this item's assets.
 
+        Args:
+            media_type: If set, filter the assets such that only those with a
+                matching ``media_type`` are returned.
+            role: If set, filter the assets such that only those with a matching
+                ``role`` are returned.
+
         Returns:
-            Dict[str, Asset]: A copy of the dictionary of this item's assets.
+            Dict[str, Asset]: A dictionary of assets that match ``media_type``
+                and/or ``role`` if set or else all of this item's assets.
         """
-        return dict(self.assets.items())
+        if media_type is None and role is None:
+            return dict(self.assets.items())
+        assets = dict()
+        for key, asset in self.assets.items():
+            if (media_type is None or asset.media_type == media_type) and (
+                role is None or asset.has_role(role)
+            ):
+                assets[key] = asset
+        return assets
 
     def add_asset(self, key: str, asset: Asset) -> None:
         """Adds an Asset to this item.

tests/data-files/collections/with-assets.json

@@ -30,7 +30,10 @@
     "thumbnail": {
       "href": "http://example.com/thumbnail.png",
       "title": "thumbnail",
-      "media_type": "image/png"
+      "type": "image/png",
+      "roles": [
+        "thumbnail"
+      ]
     }
   },
   "extent": {

tests/data-files/item/sample-item.json

@@ -54,11 +54,20 @@
     "analytic": {
       "href": "http://cool-sat.com/catalog/CS3-20160503_132130_04/analytic.tif",
       "title": "4-Band Analytic",
-      "product": "http://cool-sat.com/catalog/products/analytic.json"
+      "product": "http://cool-sat.com/catalog/products/analytic.json",
+      "type": "image/tiff; application=geotiff; profile=cloud-optimized",
+      "roles": [
+        "data",
+        "analytic"
+      ]
     },
     "thumbnail": {
       "href": "http://cool-sat.com/catalog/CS3-20160503_132130_04/thumbnail.png",
-      "title": "Thumbnail"
+      "title": "Thumbnail",
+      "type": "image/png",
+      "roles": [
+        "thumbnail"
+      ]
     }
   },
   "bbox": [

tests/test_collection.py

@@ -226,6 +226,25 @@ def test_assets(self) -> None:
         collection = pystac.Collection.from_dict(data)
         collection.validate()
 
+    def test_get_assets(self) -> None:
+        collection = pystac.Collection.from_file(
+            TestCases.get_path("data-files/collections/with-assets.json")
+        )
+
+        media_type_filter = collection.get_assets(media_type=pystac.MediaType.PNG)
+        self.assertCountEqual(media_type_filter.keys(), ["thumbnail"])
+        role_filter = collection.get_assets(role="thumbnail")
+        self.assertCountEqual(role_filter.keys(), ["thumbnail"])
+        multi_filter = collection.get_assets(
+            media_type=pystac.MediaType.PNG, role="thumbnail"
+        )
+        self.assertCountEqual(multi_filter.keys(), ["thumbnail"])
+
+        no_filter = collection.get_assets()
+        self.assertCountEqual(no_filter.keys(), ["thumbnail"])
+        no_assets = collection.get_assets(media_type=pystac.MediaType.HDF)
+        self.assertEqual(no_assets, {})
+
     def test_removing_optional_attributes(self) -> None:
         path = TestCases.get_path("data-files/collections/with-assets.json")
         with open(path, "r") as file:

tests/test_item.py

@@ -161,6 +161,25 @@ def test_null_datetime(self) -> None:
 
         null_dt_item.validate()
 
+    def test_get_assets(self) -> None:
+        item = pystac.Item.from_file(
+            TestCases.get_path("data-files/item/sample-item.json")
+        )
+
+        media_type_filter = item.get_assets(media_type=pystac.MediaType.COG)
+        self.assertCountEqual(media_type_filter.keys(), ["analytic"])
+        role_filter = item.get_assets(role="data")
+        self.assertCountEqual(role_filter.keys(), ["analytic"])
+        multi_filter = item.get_assets(
+            media_type=pystac.MediaType.PNG, role="thumbnail"
+        )
+        self.assertCountEqual(multi_filter.keys(), ["thumbnail"])
+
+        no_filter = item.get_assets()
+        self.assertCountEqual(no_filter.keys(), ["analytic", "thumbnail"])
+        no_assets = item.get_assets(media_type=pystac.MediaType.HDF)
+        self.assertEqual(no_assets, {})
+
     def test_get_set_asset_datetime(self) -> None:
         item = pystac.Item.from_file(
             TestCases.get_path("data-files/item/sample-item-asset-properties.json")