Handle extra fields in extent classes

Author: duckontheweb

pystac/collection.py

@@ -46,28 +46,41 @@ class SpatialExtent:
             array must be 2*n where n is the number of dimensions. For example, a
             2D Collection with only one bbox would be [[xmin, ymin, xmax, ymax]]
 
-    Attributes:
-        bboxes : A list of bboxes that represent the spatial
-            extent of the collection. Each bbox can be 2D or 3D. The length of the bbox
-            array must be 2*n where n is the number of dimensions. For example, a
-            2D Collection with only one bbox would be [[xmin, ymin, xmax, ymax]]
+        extra_fields : Dictionary containing additional top-level fields defined on the
+            Spatial Extent object.
     """
 
-    def __init__(self, bboxes: Union[List[List[float]], List[float]]) -> None:
+    bboxes: List[List[float]]
+    """A list of bboxes that represent the spatial
+    extent of the collection. Each bbox can be 2D or 3D. The length of the bbox
+    array must be 2*n where n is the number of dimensions. For example, a
+    2D Collection with only one bbox would be [[xmin, ymin, xmax, ymax]]"""
+
+    extra_fields: Dict[str, Any]
+    """Dictionary containing additional top-level fields defined on the Spatial
+    Extent object."""
+
+    def __init__(
+        self,
+        bboxes: Union[List[List[float]], List[float]],
+        extra_fields: Optional[Dict[str, Any]] = None,
+    ) -> None:
         # A common mistake is to pass in a single bbox instead of a list of bboxes.
         # Account for this by transforming the input in that case.
         if isinstance(bboxes, list) and isinstance(bboxes[0], float):
             self.bboxes: List[List[float]] = [cast(List[float], bboxes)]
         else:
             self.bboxes = cast(List[List[float]], bboxes)
 
+        self.extra_fields = extra_fields or {}
+
     def to_dict(self) -> Dict[str, Any]:
         """Generate a dictionary representing the JSON of this SpatialExtent.
 
         Returns:
             dict: A serialization of the SpatialExtent that can be written out as JSON.
         """
-        d = {"bbox": self.bboxes}
+        d = {"bbox": self.bboxes, **self.extra_fields}
         return d
 
     def clone(self) -> "SpatialExtent":
@@ -76,26 +89,34 @@ def clone(self) -> "SpatialExtent":
         Returns:
             SpatialExtent: The clone of this object.
         """
-        return SpatialExtent(deepcopy(self.bboxes))
+        return SpatialExtent(
+            bboxes=deepcopy(self.bboxes), extra_fields=deepcopy(self.extra_fields)
+        )
 
     @staticmethod
     def from_dict(d: Dict[str, Any]) -> "SpatialExtent":
-        """Constructs an SpatialExtent from a dict.
+        """Constructs a SpatialExtent from a dict.
 
         Returns:
             SpatialExtent: The SpatialExtent deserialized from the JSON dict.
         """
-        return SpatialExtent(bboxes=d["bbox"])
+        return SpatialExtent(
+            bboxes=d["bbox"], extra_fields={k: v for k, v in d.items() if k != "bbox"}
+        )
 
     @staticmethod
-    def from_coordinates(coordinates: List[Any]) -> "SpatialExtent":
+    def from_coordinates(
+        coordinates: List[Any], extra_fields: Optional[Dict[str, Any]] = None
+    ) -> "SpatialExtent":
         """Constructs a SpatialExtent from a set of coordinates.
 
         This method will only produce a single bbox that covers all points
         in the coordinate set.
 
         Args:
             coordinates : Coordinates to derive the bbox from.
+            extra_fields : Dictionary containing additional top-level fields defined on
+                the Spatial Extent object.
 
         Returns:
             SpatialExtent: A SpatialExtent with a single bbox that covers the
@@ -133,31 +154,40 @@ def process_coords(
                 f"Could not determine bounds from coordinate sequence {coordinates}"
             )
 
-        return SpatialExtent([[xmin, ymin, xmax, ymax]])
+        return SpatialExtent(
+            bboxes=[[xmin, ymin, xmax, ymax]], extra_fields=extra_fields
+        )
 
 
 class TemporalExtent:
     """Describes the temporal extent of a Collection.
 
     Args:
         intervals :  A list of two datetimes wrapped in a list,
-        representing the temporal extent of a Collection. Open date ranges are supported
-        by setting either the start (the first element of the interval) or the end (the
-        second element of the interval) to None.
-
-
-    Attributes:
-        intervals :  A list of two datetimes wrapped in a list,
-        representing the temporal extent of a Collection. Open date ranges are
-        represented by either the start (the first element of the interval) or the
-        end (the second element of the interval) being None.
+            representing the temporal extent of a Collection. Open date ranges are
+            supported by setting either the start (the first element of the interval)
+            or the end (the second element of the interval) to None.
 
+        extra_fields : Dictionary containing additional top-level fields defined on the
+            Temporal Extent object.
     Note:
         Datetimes are required to be in UTC.
     """
 
+    intervals: List[List[Optional[Datetime]]]
+    """A list of two datetimes wrapped in a list,
+    representing the temporal extent of a Collection. Open date ranges are
+    represented by either the start (the first element of the interval) or the
+    end (the second element of the interval) being None."""
+
+    extra_fields: Dict[str, Any]
+    """Dictionary containing additional top-level fields defined on the Temporal
+    Extent object."""
+
     def __init__(
-        self, intervals: Union[List[List[Optional[Datetime]]], List[Optional[Datetime]]]
+        self,
+        intervals: Union[List[List[Optional[Datetime]]], List[Optional[Datetime]]],
+        extra_fields: Optional[Dict[str, Any]] = None,
     ):
         # A common mistake is to pass in a single interval instead of a
         # list of intervals. Account for this by transforming the input
@@ -167,6 +197,8 @@ def __init__(
         else:
             self.intervals = cast(List[List[Optional[Datetime]]], intervals)
 
+        self.extra_fields = extra_fields or {}
+
     def to_dict(self) -> Dict[str, Any]:
         """Generate a dictionary representing the JSON of this TemporalExtent.
 
@@ -186,7 +218,7 @@ def to_dict(self) -> Dict[str, Any]:
 
             encoded_intervals.append([start, end])
 
-        d = {"interval": encoded_intervals}
+        d = {"interval": encoded_intervals, **self.extra_fields}
         return d
 
     def clone(self) -> "TemporalExtent":
@@ -195,7 +227,9 @@ def clone(self) -> "TemporalExtent":
         Returns:
             TemporalExtent: The clone of this object.
         """
-        return TemporalExtent(intervals=deepcopy(self.intervals))
+        return TemporalExtent(
+            intervals=deepcopy(self.intervals), extra_fields=deepcopy(self.extra_fields)
+        )
 
     @staticmethod
     def from_dict(d: Dict[str, Any]) -> "TemporalExtent":
@@ -215,7 +249,10 @@ def from_dict(d: Dict[str, Any]) -> "TemporalExtent":
                 end = dateutil.parser.parse(i[1])
             parsed_intervals.append([start, end])
 
-        return TemporalExtent(intervals=parsed_intervals)
+        return TemporalExtent(
+            intervals=parsed_intervals,
+            extra_fields={k: v for k, v in d.items() if k != "interval"},
+        )
 
     @staticmethod
     def from_now() -> "TemporalExtent":
@@ -236,23 +273,41 @@ class Extent:
     Args:
         spatial : Potential spatial extent covered by the collection.
         temporal : Potential temporal extent covered by the collection.
-
-    Attributes:
-        spatial : Potential spatial extent covered by the collection.
-        temporal : Potential temporal extent covered by the collection.
+        extra_fields : Dictionary containing additional top-level fields defined on the
+            Extent object.
     """
 
-    def __init__(self, spatial: SpatialExtent, temporal: TemporalExtent):
+    spatial: SpatialExtent
+    """Potential spatial extent covered by the collection."""
+
+    temporal: TemporalExtent
+    """Potential temporal extent covered by the collection."""
+
+    extra_fields: Dict[str, Any]
+    """Dictionary containing additional top-level fields defined on the Extent
+    object."""
+
+    def __init__(
+        self,
+        spatial: SpatialExtent,
+        temporal: TemporalExtent,
+        extra_fields: Optional[Dict[str, Any]] = None,
+    ):
         self.spatial = spatial
         self.temporal = temporal
+        self.extra_fields = extra_fields or {}
 
     def to_dict(self) -> Dict[str, Any]:
         """Generate a dictionary representing the JSON of this Extent.
 
         Returns:
             dict: A serialization of the Extent that can be written out as JSON.
         """
-        d = {"spatial": self.spatial.to_dict(), "temporal": self.temporal.to_dict()}
+        d = {
+            "spatial": self.spatial.to_dict(),
+            "temporal": self.temporal.to_dict(),
+            **self.extra_fields,
+        }
 
         return d
 
@@ -262,7 +317,11 @@ def clone(self) -> "Extent":
         Returns:
             Extent: The clone of this extent.
         """
-        return Extent(spatial=copy(self.spatial), temporal=copy(self.temporal))
+        return Extent(
+            spatial=copy(self.spatial),
+            temporal=copy(self.temporal),
+            extra_fields=deepcopy(self.extra_fields),
+        )
 
     @staticmethod
     def from_dict(d: Dict[str, Any]) -> "Extent":
@@ -287,16 +346,23 @@ def from_dict(d: Dict[str, Any]) -> "Extent":
             temporal_extent_dict = temporal_extent
 
         return Extent(
-            SpatialExtent.from_dict(spatial_extent_dict),
-            TemporalExtent.from_dict(temporal_extent_dict),
+            spatial=SpatialExtent.from_dict(spatial_extent_dict),
+            temporal=TemporalExtent.from_dict(temporal_extent_dict),
+            extra_fields={
+                k: v for k, v in d.items() if k not in {"spatial", "temporal"}
+            },
         )
 
     @staticmethod
-    def from_items(items: Iterable["Item_Type"]) -> "Extent":
+    def from_items(
+        items: Iterable["Item_Type"], extra_fields: Optional[Dict[str, Any]] = None
+    ) -> "Extent":
         """Create an Extent based on the datetimes and bboxes of a list of items.
 
         Args:
             items : A list of items to derive the extent from.
+            extra_fields : Optional dictionary containing additional top-level fields
+                defined on the Extent object.
 
         Returns:
             Extent: An Extent that spatially and temporally covers all of the
@@ -354,7 +420,7 @@ def from_items(items: Iterable["Item_Type"]) -> "Extent":
         )
         temporal = TemporalExtent([[start_timestamp, end_timestamp]])
 
-        return Extent(spatial, temporal)
+        return Extent(spatial=spatial, temporal=temporal, extra_fields=extra_fields)
 
 
 class Provider:

tests/test_collection.py

@@ -252,6 +252,9 @@ def test_from_invalid_dict_raises_exception(self) -> None:
 
 
 class ExtentTest(unittest.TestCase):
+    def setUp(self) -> None:
+        self.maxDiff = None
+
     def test_spatial_allows_single_bbox(self) -> None:
         temporal_extent = TemporalExtent(intervals=[[TEST_DATETIME, None]])
 
@@ -319,6 +322,49 @@ def test_from_items(self) -> None:
         self.assertEqual(interval[0], datetime(2000, 1, 1, 12, 0, 0, 0, tzinfo=tz.UTC))
         self.assertEqual(interval[1], datetime(2001, 1, 1, 12, 0, 0, 0, tzinfo=tz.UTC))
 
+    def test_to_from_dict(self) -> None:
+        spatial_dict = {
+            "bbox": [
+                [
+                    172.91173669923782,
+                    1.3438851951615003,
+                    172.95469614953714,
+                    1.3690476620161975,
+                ]
+            ],
+            "extension:field": "spatial value",
+        }
+        temporal_dict = {
+            "interval": [
+                ["2020-12-11T22:38:32.125000Z", "2020-12-14T18:02:31.437000Z"]
+            ],
+            "extension:field": "temporal value",
+        }
+        extent_dict = {
+            "spatial": spatial_dict,
+            "temporal": temporal_dict,
+            "extension:field": "extent value",
+        }
+        expected_extent_extra_fields = {
+            "extension:field": extent_dict["extension:field"],
+        }
+        expected_spatial_extra_fields = {
+            "extension:field": spatial_dict["extension:field"],
+        }
+        expected_temporal_extra_fields = {
+            "extension:field": temporal_dict["extension:field"],
+        }
+
+        extent = Extent.from_dict(extent_dict)
+
+        self.assertDictEqual(expected_extent_extra_fields, extent.extra_fields)
+        self.assertDictEqual(expected_spatial_extra_fields, extent.spatial.extra_fields)
+        self.assertDictEqual(
+            expected_temporal_extra_fields, extent.temporal.extra_fields
+        )
+
+        self.assertDictEqual(extent_dict, extent.to_dict())
+
 
 class CollectionSubClassTest(unittest.TestCase):
     """This tests cases related to creating classes inheriting from pystac.Catalog to