fix concat with variable or dataarray as dim (#6387)

Propagate attrs.

Author: benbovy

xarray/core/concat.py

@@ -429,6 +429,7 @@ def _dataset_concat(
     """
     Concatenate a sequence of datasets along a new or existing dimension
     """
+    from .dataarray import DataArray
     from .dataset import Dataset
 
     datasets = list(datasets)
@@ -438,6 +439,13 @@ def _dataset_concat(
             "The elements in the input list need to be either all 'Dataset's or all 'DataArray's"
         )
 
+    if isinstance(dim, DataArray):
+        dim_var = dim.variable
+    elif isinstance(dim, Variable):
+        dim_var = dim
+    else:
+        dim_var = None
+
     dim, index = _calc_concat_dim_index(dim)
 
     # Make sure we're working on a copy (we'll be loading variables)
@@ -582,7 +590,11 @@ def get_indexes(name):
 
     if index is not None:
         # add concat index / coordinate last to ensure that its in the final Dataset
-        result[dim] = index.create_variables()[dim]
+        if dim_var is not None:
+            index_vars = index.create_variables({dim: dim_var})
+        else:
+            index_vars = index.create_variables()
+        result[dim] = index_vars[dim]
         result_indexes[dim] = index
 
     # TODO: add indexes at Dataset creation (when it is supported)

xarray/tests/test_concat.py

@@ -459,8 +459,15 @@ def test_concat_do_not_promote(self) -> None:
 
     def test_concat_dim_is_variable(self) -> None:
         objs = [Dataset({"x": 0}), Dataset({"x": 1})]
-        coord = Variable("y", [3, 4])
-        expected = Dataset({"x": ("y", [0, 1]), "y": [3, 4]})
+        coord = Variable("y", [3, 4], attrs={"foo": "bar"})
+        expected = Dataset({"x": ("y", [0, 1]), "y": coord})
+        actual = concat(objs, coord)
+        assert_identical(actual, expected)
+
+    def test_concat_dim_is_dataarray(self) -> None:
+        objs = [Dataset({"x": 0}), Dataset({"x": 1})]
+        coord = DataArray([3, 4], dims="y", attrs={"foo": "bar"})
+        expected = Dataset({"x": ("y", [0, 1]), "y": coord})
         actual = concat(objs, coord)
         assert_identical(actual, expected)