Add variables argument to the DatacubeExtension apply method (#782)

* let apply set variables

* add test for datacube apply

* pacify mypy

* add to changelog

Author: itcarroll

CHANGELOG.md

@@ -16,6 +16,7 @@
 ### Fixed
 
 - "How to create STAC catalogs" tutorial ([#775](https://github.com/stac-utils/pystac/pull/775))
+- Add a `variables` argument, to accompany `dimensions`, for the `apply` method of stac objects extended with datacube ([#782](https://github.com/stac-utils/pystac/pull/782))
 
 ## [v1.4.0]
 

pystac/extensions/datacube.py

@@ -411,6 +411,12 @@ class VariableType(StringEnum):
 
 
 class Variable:
+    """Object representing a variable in the datacube. The dimensions field lists
+    zero or more :stac-ext:`Datacube Dimension Object <datacube#dimension-object>`
+    instances. See the :stac-ext:`Datacube Variable Object
+    <datacube#variable-object>` docs for details.
+    """
+
     properties: Dict[str, Any]
 
     def __init__(self, properties: Dict[str, Any]) -> None:
@@ -522,15 +528,22 @@ class DatacubeExtension(
        >>> dc_ext = DatacubeExtension.ext(item)
     """
 
-    def apply(self, dimensions: Dict[str, Dimension]) -> None:
+    def apply(
+        self,
+        dimensions: Dict[str, Dimension],
+        variables: Optional[Dict[str, Variable]] = None,
+    ) -> None:
         """Applies label extension properties to the extended
         :class:`~pystac.Collection`, :class:`~pystac.Item` or :class:`~pystac.Asset`.
 
         Args:
             dimensions : Dictionary mapping dimension name to a :class:`Dimension`
                 object.
+            variables : Dictionary mapping variable name to a :class:`Variable`
+                object.
         """
         self.dimensions = dimensions
+        self.variables = variables
 
     @property
     def dimensions(self) -> Dict[str, Dimension]:

tests/extensions/test_datacube.py

@@ -87,3 +87,16 @@ def test_set_variables(self) -> None:
         self.assertEqual(
             item.properties["cube:variables"], {"temp": new_variable.to_dict()}
         )
+
+    def test_apply_variables(self) -> None:
+        item = pystac.Item.from_file(self.example_uri)
+        cube = DatacubeExtension.ext(item)
+        variables = cube.variables
+        assert variables is not None
+        key, value = variables.popitem()
+        target = value.to_dict()
+        cube.variables = None
+        cube.apply(dimensions={}, variables={key: value})
+        variables = cube.variables
+        assert variables is not None
+        self.assertEqual(target, cube.variables[key].to_dict())