Allow pd.Index arguments as 'dim' in xray.concat

Author: shoyer

xray/core/dataset.py

@@ -1341,11 +1341,20 @@ def concat(cls, *args, **kwargs):
     @classmethod
     def _concat(cls, datasets, dim='concat_dim', indexers=None,
                 mode='different', concat_over=None, compat='equals'):
+        from .dataarray import DataArray
+
         _assert_compat_valid(compat)
 
         # don't bother trying to work with datasets as a generator instead of a
         # list; the gains would be minimal
         datasets = list(map(as_dataset, datasets))
+
+        if not isinstance(dim, basestring) and not hasattr(dim, 'dims'):
+            # dim is not a DataArray or Coordinate
+            dim_name = getattr(dim, 'name', None)
+            if dim_name is None:
+                dim_name = 'concat_dim'
+            dim = DataArray(dim, dims=dim_name, name=dim_name)
         dim_name = getattr(dim, 'name', dim)
 
         # figure out variables to concatenate over

xray/test/test_dataarray.py

@@ -826,6 +826,17 @@ def test_concat(self):
         grouped = [g for _, g in foo.groupby('x')]
         stacked = concat(grouped, self.ds['x'])
         self.assertDataArrayIdentical(foo, stacked)
+        # with an index as the 'dim' argument
+        stacked = concat(grouped, self.ds.indexes['x'])
+        self.assertDataArrayIdentical(foo, stacked)
+
+        actual = concat([foo[0], foo[1]], pd.Index([0, 1])).reset_coords(drop=True)
+        expected = foo[:2].rename({'x': 'concat_dim'})
+        self.assertDataArrayIdentical(expected, actual)
+
+        actual = concat([foo[0], foo[1]], [0, 1]).reset_coords(drop=True)
+        expected = foo[:2].rename({'x': 'concat_dim'})
+        self.assertDataArrayIdentical(expected, actual)
 
         with self.assertRaisesRegexp(ValueError, 'not identical'):
             concat([foo, bar], compat='identical')

xray/test/test_dataset.py

@@ -938,7 +938,7 @@ def rectify_dim_order(dataset):
         expected['dim1'] = dim
         self.assertDatasetIdentical(expected, concat(datasets, dim))
 
-        # TODO: factor this into several distinct tests
+    def test_concat_errors(self):
         data = create_test_data()
         split_data = [data.isel(dim1=slice(10)),
                       data.isel(dim1=slice(10, None))]