diff --git a/src/pystac/__init__.py b/src/pystac/__init__.py index 0ef340657..a19849bb4 100644 --- a/src/pystac/__init__.py +++ b/src/pystac/__init__.py @@ -3,6 +3,7 @@ from typing_extensions import deprecated from .asset import Asset +from .band import Band from .catalog import Catalog from .collection import ( Collection, @@ -61,6 +62,7 @@ def get_stac_version() -> str: __all__ = [ "Asset", + "Band", "Catalog", "Collection", "CommonMetadata", diff --git a/src/pystac/asset.py b/src/pystac/asset.py index d12f858d7..d974c876f 100644 --- a/src/pystac/asset.py +++ b/src/pystac/asset.py @@ -5,6 +5,7 @@ from typing_extensions import deprecated +from .band import Band from .media_type import MediaType from .utils import is_absolute_href, make_absolute_href, make_relative_href from .writer import Writer @@ -20,12 +21,17 @@ def __init__( description: str | None = None, type: str | None = None, roles: list[str] | None = None, + bands: list[Band] | list[dict[str, Any]] | None = None, **kwargs: Any, ): self.title: str | None = title self.description: str | None = description self.type: str | None = type self.roles: list[str] | None = roles + if bands is not None: + self.bands: list[Band] | None = [Band.try_from(band) for band in bands] + else: + self.bands = None self.extra_fields: dict[str, Any] = kwargs @classmethod @@ -54,6 +60,8 @@ def to_dict(self) -> dict[str, Any]: data["type"] = self.type if self.roles is not None: data["roles"] = self.roles + if self.bands is not None: + data["bands"] = [band.to_dict() for band in self.bands] return data @@ -65,10 +73,16 @@ def __init__( description: str | None = None, type: str | None = None, roles: list[str] | None = None, + bands: list[Band] | list[dict[str, Any]] | None = None, **kwargs: Any, ): super().__init__( - title=title, description=description, type=type, roles=roles, **kwargs + title=title, + description=description, + type=type, + roles=roles, + bands=bands, + **kwargs, ) self.href: str = href diff --git a/src/pystac/band.py b/src/pystac/band.py new file mode 100644 index 000000000..837161600 --- /dev/null +++ b/src/pystac/band.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import copy +from typing import Any, override + + +class Band: + def __init__( + self, + name: str | None = None, + **kwargs: Any, + ): + self.name: str | None = name + self.extra_fields: dict[str, Any] = kwargs + + def __getitem__(self, key: str) -> Any: + if hasattr(self, key): + return getattr(self, key) + else: + return self.extra_fields[key] + + def __setitem__(self, key: str, value: Any) -> None: + if hasattr(self, key): + setattr(self, key, value) + else: + self.extra_fields[key] = value + + def __contains__(self, key: str) -> bool: + return key == "name" or key in self.extra_fields + + def __deepcopy__(self, memo: dict[int, Any]) -> Band: + return Band( + name=self.name, + **copy.deepcopy(self.extra_fields, memo), + ) + + @override + def __eq__(self, other: Any) -> bool: + if not isinstance(other, Band): + return NotImplemented + return self.to_dict() == other.to_dict() + + @override + def __repr__(self) -> str: + return f"Band(name={self.name!r})" + + @classmethod + def try_from(cls, data: dict[str, Any] | Band) -> Band: + if isinstance(data, Band): + return data + else: + return cls(**data) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Band: + return cls(**data) + + def to_dict(self) -> dict[str, Any]: + data = copy.deepcopy(self.extra_fields) + if self.name: + data["name"] = self.name + return data diff --git a/src/pystac/collection.py b/src/pystac/collection.py index c06b68882..df5578ebb 100644 --- a/src/pystac/collection.py +++ b/src/pystac/collection.py @@ -7,6 +7,7 @@ from typing_extensions import deprecated from .asset import Asset, Assets, ItemAsset +from .band import Band from .constants import DEFAULT_LICENSE, DEFAULT_STAC_VERSION, STAC_OBJECT_TYPE from .container import Container from .link import Link @@ -46,6 +47,7 @@ def __init__( stac_version: str = DEFAULT_STAC_VERSION, stac_extensions: list[str] | None = None, links: list[Link] | list[dict[str, Any]] | None = None, + bands: list[Band] | list[dict[str, Any]] | None = None, **kwargs: Any, ) -> None: super().__init__(type, id, stac_version, stac_extensions, links, **kwargs) @@ -58,6 +60,10 @@ def __init__( if providers is not None else None ) + if bands is not None: + self.bands: list[Band] | None = [Band.try_from(band) for band in bands] + else: + self.bands = None self.extent: Extent = Extent.try_from(extent) self.summaries: dict[str, Any] | None = summaries self.assets: dict[str, Asset] = ( @@ -108,6 +114,8 @@ def to_dict( data["extent"] = self.extent.to_dict() if self.summaries: data["summaries"] = self.summaries + if self.bands is not None: + data["bands"] = [band.to_dict() for band in self.bands] if self.assets: data["assets"] = { key: asset.to_dict() for key, asset in self.assets.items() diff --git a/src/pystac/item.py b/src/pystac/item.py index 3b38aed1c..19118db5b 100644 --- a/src/pystac/item.py +++ b/src/pystac/item.py @@ -10,6 +10,7 @@ from pystac.rel_type import RelType from .asset import Asset, Assets +from .band import Band from .constants import DEFAULT_STAC_VERSION, STAC_OBJECT_TYPE from .container import Container from .geo_interface import GeoInterface @@ -218,7 +219,12 @@ def __geo_interface__(self) -> dict[str, Any]: class Properties: - def __init__(self, datetime: dt.datetime | str | None = None, **kwargs: Any): + def __init__( + self, + datetime: dt.datetime | str | None = None, + bands: list[Band] | list[dict[str, Any]] | None = None, + **kwargs: Any, + ): if isinstance(datetime, str): self.datetime: dt.datetime | None = str_to_datetime(datetime) elif isinstance(datetime, dt.datetime): @@ -226,7 +232,10 @@ def __init__(self, datetime: dt.datetime | str | None = None, **kwargs: Any): else: # TODO check for start and end datetime self.datetime = dt.datetime.now(tz=dt.UTC) - + if bands is not None: + self.bands: list[Band] | None = [Band.try_from(band) for band in bands] + else: + self.bands = None self.extra_fields: dict[str, Any] = kwargs def __getitem__(self, key: str) -> Any: @@ -236,11 +245,12 @@ def __setitem__(self, name: str, value: Any, /) -> None: self.extra_fields[name] = value def __contains__(self, key: str) -> bool: - return key == "datetime" or key in self.extra_fields + return key in ("datetime", "bands") or key in self.extra_fields def __deepcopy__(self, memo: dict[int, Any]) -> Properties: return Properties( datetime=self.datetime, + bands=self.bands, **copy.deepcopy(self.extra_fields, memo), ) @@ -263,4 +273,6 @@ def from_dict(cls, data: dict[str, Any]) -> Properties: def to_dict(self) -> dict[str, Any]: data = copy.deepcopy(self.extra_fields) data["datetime"] = datetime_to_str(self.datetime) if self.datetime else None + if self.bands is not None: + data["bands"] = [band.to_dict() for band in self.bands] return data diff --git a/tests/test_band.py b/tests/test_band.py new file mode 100644 index 000000000..e2f25b8b9 --- /dev/null +++ b/tests/test_band.py @@ -0,0 +1,94 @@ +import pytest + +from pystac import Asset, Band, Collection, Item + + +@pytest.mark.parametrize( + "data", [{"name": "B01"}, {}, {"name": "B01", "eo:common_name": "blue"}] +) +def test_roundtrip_to_dict(data: dict[str, str]) -> None: + band = Band(**data) + assert band.to_dict() == data + + +def test_get_band_fields() -> None: + band = Band(name="foo", field="one") + assert band["name"] == "foo" + assert band["field"] == "one" + assert "field" in band + assert band.extra_fields["field"] == "one" + + +def test_set_band_fields() -> None: + band = Band() + band["field"] = "two" + assert "field" in band + assert band["field"] == "two" + assert band.extra_fields["field"] == "two" + + +def test_set_band_name() -> None: + band = Band() + band.name = "B01" + assert band.name == "B01" + assert band["name"] == "B01" + + band["name"] = "B02" + assert band.name == "B02" + assert band["name"] == "B02" + + +def test_band_repr() -> None: + band = Band(name="foo", field="one") + assert str(band) == "Band(name='foo')" + + +def test_set_bands_on_asset() -> None: + asset = Asset(href="foo", bands=[Band(name="B01")]) + assert asset.bands == [Band(name="B01")] + + +def test_set_bands_on_asset_as_dict() -> None: + asset = Asset(href="foo", bands=[{"name": "B01"}]) + assert asset.bands == [Band(name="B01")] + + +def test_asset_to_dict_includes_bands() -> None: + asset = Asset(href="foo", bands=[Band("B01")]) + asset_dict = asset.to_dict() + assert "bands" in asset_dict + assert asset_dict["bands"] == [{"name": "B01"}] + + +def test_set_bands_on_item() -> None: + item = Item(id="foo", properties={"bands": [Band("B01")]}) + assert item.properties.bands == [Band("B01")] + + +def test_set_bands_on_item_as_dict() -> None: + item = Item(id="foo", properties={"bands": [{"name": "B01"}]}) + assert item.properties.bands == [Band("B01")] + + +def test_item_to_dict_includes_bands() -> None: + item = Item(id="foo", properties={"bands": [Band("B01")]}) + properties_dict = item.to_dict()["properties"] + assert "bands" in properties_dict + assert properties_dict["bands"] == [{"name": "B01"}] + + +def test_set_bands_on_collection() -> None: + collection = Collection(id="foo", description="foo", bands=[Band("B01")]) + assert collection.bands == [Band("B01")] + + +def test_set_bands_on_collection_as_dict() -> None: + collection = Collection(id="foo", description="foo", bands=[{"name": "B01"}]) + assert collection.bands == [Band("B01")] + + +def test_collection_to_dict_includes_bands() -> None: + collection = Collection(id="foo", description="foo", bands=[Band("B01")]) + collection_dict = collection.to_dict() + assert "bands" in collection_dict + assert collection_dict["bands"] == [{"name": "B01"}]