diff --git a/src/stactools/ecmwf_forecast/stac.py b/src/stactools/ecmwf_forecast/stac.py index 9f1f0b2..c5a628f 100644 --- a/src/stactools/ecmwf_forecast/stac.py +++ b/src/stactools/ecmwf_forecast/stac.py @@ -7,7 +7,7 @@ import operator import pathlib import re -from typing import Any +from typing import Any, Optional import fsspec import pystac @@ -48,6 +48,7 @@ class Parts: format: str filename: str split_by_step: bool = False + resolution: Optional[str] = None @property def datetime(self): @@ -57,7 +58,7 @@ def datetime(self): return self.reference_datetime @classmethod - def from_filename(cls, filename: str, split_by_step=False) -> "Parts": + def from_filename(cls, filename: str, split_by_step=False, resolution: Optional[str] = None) -> "Parts": name = pathlib.Path(filename).name m = xpr.match(name) if not m: @@ -67,6 +68,8 @@ def from_filename(cls, filename: str, split_by_step=False) -> "Parts": d["reference_datetime"] = datetime.datetime.strptime( d["reference_datetime"], "%Y%m%d%H" ) # type: ignore + if resolution: + d["resolution"] = resolution # error: Argument 1 to "Parts" has incompatible type # "**Dict[str, Union[str, Any]]"; expected "datetime" return cls(**d, split_by_step=split_by_step) # type: ignore @@ -81,6 +84,8 @@ def item_id(self) -> str: ] if self.split_by_step: parts.append(self.step) + if self.resolution is not None: + parts.append(self.resolution) return "-".join(parts) @property @@ -279,7 +284,7 @@ def group_assets(asset_hrefs: list[str], key=item_key): return grouped -def create_item(asset_hrefs: list[str], split_by_step=False) -> Item: +def create_item(asset_hrefs: list[str], split_by_step=False, resolution: Optional[str] = None) -> Item: """ Create an item for the hrefs. @@ -295,7 +300,7 @@ def create_item(asset_hrefs: list[str], split_by_step=False) -> Item: pystac.Item """ parts = [ - Parts.from_filename(href, split_by_step=split_by_step) for href in asset_hrefs + Parts.from_filename(href, split_by_step=split_by_step, resolution=resolution) for href in asset_hrefs ] return _create_item_from_parts(parts, split_by_step=split_by_step) @@ -362,6 +367,7 @@ def _create_item_from_parts(parts: list[Parts], split_by_step=False) -> Item: item.properties["ecmwf:reference_datetime"] = ( part.reference_datetime.isoformat() + "Z" ) + item.properties["ecmwf:resolution"] = part.resolution item.properties["ecmwf:forecast_datetime"] = ( part.forecast_datetime.isoformat() + "Z" ) diff --git a/tests/test_stac.py b/tests/test_stac.py index ffedaca..32e2650 100644 --- a/tests/test_stac.py +++ b/tests/test_stac.py @@ -130,6 +130,30 @@ def test_split_by_parts(): assert i1.id.endswith("3h") +@pytest.mark.parametrize( + ["files", "resolution"], + [ + ( + ( + "ecmwf/20220222/00z/0p4-beta/enfo/20220222000000-0h-enfo-ef.grib2", + "ecmwf/20220222/00z/0p4-beta/enfo/20220222000000-0h-enfo-ef.index", + ), + "0.40", + ), + ( + ( + "ecmwf/20220222/00z/0p25/enfo/20220222000000-0h-enfo-ef.grib2", + "ecmwf/20220222/00z/0p25/enfo/20220222000000-0h-enfo-ef.index", + ), + "0.25", + ), + ], +) +def test_resolution(files, resolution): + result = stac.create_item(files, resolution=resolution) + assert result.properties["ecmwf:resolution"] == resolution + assert result.id == f"ecmwf-2022-02-22T00-enfo-ef-{resolution}" + def test_item_assets(): collection = stac.create_collection()