feat: use regex to select band in StacAssets (#38)

CS-SI · Nov 14, 2023 · 4ee07f0 · 4ee07f0
1 parent 5a31f7a
commit 4ee07f0
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 11 deletions.
diff --git a/eodag_cube/api/product/drivers/stac_assets.py b/eodag_cube/api/product/drivers/stac_assets.py
@@ -15,6 +15,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import re
+
 from eodag.api.product.drivers.base import DatasetDriver
 from eodag.utils.exceptions import AddressNotFound
 
@@ -28,11 +30,47 @@ def get_data_address(self, eo_product, band):
         See :func:`~eodag.api.product.drivers.base.DatasetDriver.get_data_address` to get help on the formal
         parameters.
         """
-        if band in eo_product.assets:
-            return eo_product.assets[band]["href"]
-        elif band.upper() in [b.upper() for b in eo_product.assets.keys()]:
-            for k, v in eo_product.assets.items():
-                if k.upper() == band.upper():
-                    return v["href"]
+        error_message = ""
+
+        # try using exact
+        p = re.compile(rf"^{band}$", re.IGNORECASE)
+        matching_keys = [
+            s
+            for s in eo_product.assets.keys()
+            if (
+                (
+                    "roles" in eo_product.assets[s]
+                    and "data" in eo_product.assets[s]["roles"]
+                )
+                or ("roles" not in eo_product.assets[s])
+            )
+            and p.match(s)
+        ]
+        if len(matching_keys) == 1:
+            return eo_product.assets[matching_keys[0]]["href"]
+        else:
+            error_message += (
+                rf"{len(matching_keys)} assets keys found matching {p} AND "
+            )
 
-        raise AddressNotFound
+            # try to find keys containing given band
+            p = re.compile(rf"^.*{band}.*$", re.IGNORECASE)
+            matching_keys = [
+                s
+                for s in eo_product.assets.keys()
+                if (
+                    (
+                        "roles" in eo_product.assets[s]
+                        and "data" in eo_product.assets[s]["roles"]
+                    )
+                    or ("roles" not in eo_product.assets[s])
+                )
+                and p.match(s)
+            ]
+            if len(matching_keys) == 1:
+                return eo_product.assets[matching_keys[0]]["href"]
+            else:
+                raise AddressNotFound(
+                    rf"Please adapt given band parameter ('{band}') to match only one asset: {error_message}"
+                    rf"{len(matching_keys)} assets keys found matching {p}"
+                )
diff --git a/tests/units/test_eoproduct_driver_stac_assets.py b/tests/units/test_eoproduct_driver_stac_assets.py
@@ -35,16 +35,39 @@ def setUp(self):
             "S2A_MSIL1C_20180101T105441_N0206_R051_T31TDH_20180101T124911.SAFE",
         )
         self.product.assets = {
-            "B01": {
-                "title": "Band 1 (coastal)",
+            "T31TDH_20180101T124911_B01.jp2": {
+                "title": "Band 1",
                 "type": "image/jp2",
+                "roles": ["data"],
                 "href": os.path.join(
                     TEST_RESOURCES_PATH,
                     "products",
                     "S2A_MSIL1C_20180101T105441_N0206_R051_T31TDH_20180101T124911.SAFE/"
                     + "GRANULE/L1C_T31TDH_A013204_20180101T105435/IMG_DATA/T31TDH_20180101T105441_B01.jp2",
                 ),
-            }
+            },
+            "T31TDH_20180101T124911_B03.jp2": {
+                "title": "Band 3",
+                "type": "image/jp2",
+                "roles": ["data"],
+                "href": os.path.join(
+                    TEST_RESOURCES_PATH,
+                    "products",
+                    "S2A_MSIL1C_20180101T105441_N0206_R051_T31TDH_20180101T124911.SAFE/"
+                    + "GRANULE/L1C_T31TDH_A013204_20180101T105435/IMG_DATA/T31TDH_20180101T105441_B03.jp2",
+                ),
+            },
+            "T31TDH_20180101T124911_B01.json": {
+                "title": "Band 1 metadata",
+                "type": "image/jp2",
+                "roles": ["metadata"],
+                "href": os.path.join(
+                    TEST_RESOURCES_PATH,
+                    "products",
+                    "S2A_MSIL1C_20180101T105441_N0206_R051_T31TDH_20180101T124911.SAFE/"
+                    + "GRANULE/L1C_T31TDH_A013204_20180101T105435/IMG_DATA/T31TDH_20180101T105441_B01.json",
+                ),
+            },
         }
         self.stac_assets_driver = StacAssets()
 
@@ -54,11 +77,22 @@ def test_driver_get_local_dataset_address_bad_band(self):
             driver = StacAssets()
             band = "B02"
             self.assertRaises(AddressNotFound, driver.get_data_address, product, band)
+            band = "B0"
+            self.assertRaises(AddressNotFound, driver.get_data_address, product, band)
 
     def test_driver_get_local_dataset_address_ok(self):
         """Driver returns a good address for an existing band"""
         with self._filesystem_product() as product:
-            band = "B01"
+            band = "b01"
+            address = self.stac_assets_driver.get_data_address(product, band)
+            self.assertEqual(address, self.local_band_file)
+            band = "t31tdh_20180101t124911_b01"
+            address = self.stac_assets_driver.get_data_address(product, band)
+            self.assertEqual(address, self.local_band_file)
+            band = "T31TDH_20180101T124911_B01.jp2"
+            address = self.stac_assets_driver.get_data_address(product, band)
+            self.assertEqual(address, self.local_band_file)
+            band = "T31TDH.*B01.*"
             address = self.stac_assets_driver.get_data_address(product, band)
             self.assertEqual(address, self.local_band_file)