Merge pull request #65 from nsidc/51-consistent-bbox

More flexible bounding box model

Author: trey-stafford

src/nsidc/iceflow/data/models.py

@@ -235,6 +235,55 @@ class BoundingBox(pydantic.BaseModel):
     upper_right_lon: float
     upper_right_lat: float
 
+    def __init__(self, *args, **kwargs):
+        """Accept either named args, one arg for each coord, or an iterable of
+        `(lower_left_lon, lower_left_lat, upper_right_lon, upper_right_lat)`."""
+        if args:
+            if len(args) == 1 and isinstance(args[0], list | tuple):
+                # The first arg should be treated like a tuple of
+                # (lower_left_lon, lower_left_lat, upper_right_lon,
+                # upper_right_lat)
+                args = tuple(args[0])
+            # Each arg should be treated as one of (lower_left_lon,
+            # lower_left_lat, upper_right_lon, upper_right_lat).
+            if len(args) != 4:
+                raise ValueError(
+                    "Expected four values for bounding box:"
+                    " (lower_left_lon, lower_left_lat, upper_right_lon, upper_right_lat)"
+                )
+            # Set kwargs if args are given.
+            kwargs = {
+                "lower_left_lon": args[0],
+                "lower_left_lat": args[1],
+                "upper_right_lon": args[2],
+                "upper_right_lat": args[3],
+            }
+
+        # Initialize the model.
+        super().__init__(**kwargs)
+
+    def __iter__(self):
+        """Return bounding box as a iter (list/tuple)."""
+        return iter(
+            (
+                self.lower_left_lon,
+                self.lower_left_lat,
+                self.upper_right_lon,
+                self.upper_right_lat,
+            )
+        )
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            return list(self.__iter__())[idx]
+        elif isinstance(idx, str):
+            return getattr(self, idx)
+        else:
+            raise TypeError(
+                "Getitem on BoundingBox must be int (e.g. 0)"
+                " or str (e.g., 'lower_left_lon')."
+            )
+
 
 ALL_DATASETS: list[Dataset] = [
     ILATM1BDataset(version="1"),

tests/unit/test_data_models.py

@@ -4,7 +4,7 @@
 import pandera as pa
 import pytest
 
-from nsidc.iceflow.data.models import IceflowDataFrame
+from nsidc.iceflow.data.models import BoundingBox, IceflowDataFrame
 
 _mock_bad_df = pd.DataFrame(
     {
@@ -45,3 +45,55 @@ def _pa_check_out(df: pd.DataFrame) -> IceflowDataFrame:
 def test_pa_check_out():
     with pytest.raises(pa.errors.SchemaError):
         _pa_check_out(_mock_bad_df)
+
+
+def test_bounding_box():
+    bbox_with_kwargs = BoundingBox(
+        lower_left_lon=-103.125559,
+        lower_left_lat=-75.180563,
+        upper_right_lon=-102.677327,
+        upper_right_lat=-74.798063,
+    )
+
+    bbox_with_args = BoundingBox(
+        -103.125559,
+        -75.180563,
+        -102.677327,
+        -74.798063,
+    )
+
+    bbox_with_list = BoundingBox(
+        [
+            -103.125559,
+            -75.180563,
+            -102.677327,
+            -74.798063,
+        ]
+    )
+
+    assert bbox_with_kwargs == bbox_with_args
+    assert bbox_with_kwargs == bbox_with_list
+
+    bbox = bbox_with_kwargs
+
+    # Access via named values
+    assert bbox.lower_left_lon == -103.125559
+    assert bbox.lower_left_lat == -75.180563
+    assert bbox.upper_right_lon == -102.677327
+    assert bbox.upper_right_lat == -74.798063
+
+    # Access as a list/tuple
+    assert list(bbox) == [-103.125559, -75.180563, -102.677327, -74.798063]
+    assert tuple(bbox) == (-103.125559, -75.180563, -102.677327, -74.798063)
+
+    # Access via index
+    assert bbox[0] == -103.125559
+    assert bbox[1] == -75.180563
+    assert bbox[2] == -102.677327
+    assert bbox[3] == -74.798063
+
+    # Access via name
+    assert bbox["lower_left_lon"] == -103.125559
+    assert bbox["lower_left_lat"] == -75.180563
+    assert bbox["upper_right_lon"] == -102.677327
+    assert bbox["upper_right_lat"] == -74.798063