Support stack (#192)

Author: dcherian

cf_xarray/accessor.py

@@ -428,7 +428,6 @@ def wrapper(obj: Union[DataArray, Dataset], key: str):
     "dim": (_get_dims,),
     "dims": (_get_dims,),  # transpose
     "drop_dims": (_get_dims,),  # drop_dims
-    "dimensions": (_get_dims,),  # stack
     "dims_dict": (_get_dims,),  # swap_dims, rename_dims
     "shifts": (_get_dims,),  # shift, roll
     "pad_width": (_get_dims,),  # shift, roll
@@ -1435,6 +1434,23 @@ def drop(self, *args, **kwargs):
             "Please use .cf.drop_vars or .cf.drop_sel as appropriate."
         )
 
+    def stack(self, dimensions=None, **dimensions_kwargs):
+        # stack needs to rewrite the _values_ of a dict
+        # our other machinery rewrites the _keys_ of a dict
+        # This seems somewhat rare, so do it explicitly for now
+
+        if dimensions is None:
+            dimensions = dimensions_kwargs
+        for key, values in dimensions.items():
+            updates = [
+                apply_mapper(
+                    (_single(_get_dims),), self._obj, v, error=True, default=[v]
+                )
+                for v in values
+            ]
+            dimensions.update({key: tuple(itertools.chain(*updates))})
+        return self._obj.stack(dimensions)
+
 
 @xr.register_dataset_accessor("cf")
 class CFDatasetAccessor(CFAccessor):

cf_xarray/tests/test_accessor.py

@@ -1041,3 +1041,13 @@ def test_groupby_special_ops():
     expected = grouped - grouped.mean()
     actual = grouped - cfgrouped.mean()
     assert_identical(expected, actual)
+
+
+@pytest.mark.parametrize("obj", objects)
+def test_stack(obj):
+    expected = obj.stack(latlon=["lat", "lon"])
+    actual = obj.cf.stack(latlon=["latitude", "longitude"])
+    assert_identical(expected, actual)
+
+    actual = obj.cf.stack({"latlon": ["latitude", "longitude"]})
+    assert_identical(expected, actual)

ci/doc.yml

@@ -6,6 +6,7 @@ dependencies:
   - python=3.8
   - matplotlib-base
   - netcdf4
+  - pooch
   - xarray
   - sphinx
   - nbsphinx

ci/environment.yml

@@ -9,4 +9,5 @@ dependencies:
   - matplotlib-base
   - netcdf4
   - pandas
+  - pooch
   - xarray

ci/upstream-dev-env.yml

@@ -9,5 +9,6 @@ dependencies:
   - matplotlib-base
   - netcdf4
   - pandas
+  - pooch
   - pip:
     - git+https://github.com/pydata/xarray