Repository: pydata/xarray

Issue:
Assigning several variables at once in a dataset
Hi,

I cannot tell if I'm doing something illegitimate. I have a hard time to understand the error I'm facing:

```python
import numpy as np
import xarray as xr

ds = xr.Dataset(
    {'varA': (['x', 'time'], np.random.rand(2, 2)),
     'varB': (['x', 'time'], np.random.rand(2, 2)),
     'varC': (['x', 'time'], np.random.rand(2, 2))},
    coords={'x': [0, 1.], 'time': [0, 1.]})

print(ds)

# I want to center 'varA' and 'varB' by substracting their mean values along the time dimension:
ds[['varA', 'varB']] -= ds[['varA', 'varB']].mean(dim='time')
```

Produces the output:
```
<xarray.Dataset>
Dimensions:  (time: 2, x: 2)
Coordinates:
  * x        (x) float64 0.0 1.0
  * time     (time) float64 0.0 1.0
Data variables:
    varA     (x, time) float64 0.7737 0.2608 0.009832 0.8498
    varC     (x, time) float64 0.9732 0.6492 0.5118 0.5792
    varB     (x, time) float64 0.04973 0.4012 0.01248 0.8716
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-36-aa87123fb25d> in <module>()
      7 print(ds)
      8 
----> 9 ds[['varA', 'varB']] -= ds[['varA', 'varB']].mean(dim='time')

C:\Users\poimat04\AppData\Local\Continuum\Miniconda2\envs\AxialLLRCA\lib\site-packages\xarray\core\dataset.pyc in __setitem__(self, key, value)
    536             raise NotImplementedError('cannot yet use a dictionary as a key '
    537                                       'to set Dataset values')
--> 538         self.update({key: value})
    539 
    540     def __delitem__(self, key):

TypeError: unhashable type: 'list'
```

Is it forbidden to assign several variables at once? I'm puzzled with this TypeError.

Thanks

Discussion:
max-sixty: I think this is a good feature — any thoughts from others? Should a dataset be a valid rvalue?

I added some comments @gcaria . Great to have you as a contributor. We could add a line or two in the docs, and whatsnew. And you're right — the tests could be filled out a bit (no need to go overboard though).
max-sixty: Great @gcaria ! And thanks for the issue link.

Unless we hear from others, we can merge in the next couple of days. Could we add a test for the ds rvalue and a whatsnew?
max-sixty: Great — @gcaria if you're OK to hit the "Commit suggestion" on those then we can merge?

Thanks for the review @mathause 
pep8speaks: Hello @gcaria! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:



* In the file [`xarray/tests/test_dataset.py`](https://github.com/pydata/xarray/blob/f6889a50d0fdd35da2197a1c8c2a7543a600b2d6/xarray/tests/test_dataset.py):

> [Line 3510:18](https://github.com/pydata/xarray/blob/f6889a50d0fdd35da2197a1c8c2a7543a600b2d6/xarray/tests/test_dataset.py#L3510): [F821](https://duckduckgo.com/?q=pep8%20F821) undefined name 'data2'
> [Line 3511:9](https://github.com/pydata/xarray/blob/f6889a50d0fdd35da2197a1c8c2a7543a600b2d6/xarray/tests/test_dataset.py#L3511): [F821](https://duckduckgo.com/?q=pep8%20F821) undefined name 'data1'
> [Line 3512:9](https://github.com/pydata/xarray/blob/f6889a50d0fdd35da2197a1c8c2a7543a600b2d6/xarray/tests/test_dataset.py#L3512): [F821](https://duckduckgo.com/?q=pep8%20F821) undefined name 'data2'
> [Line 3513:26](https://github.com/pydata/xarray/blob/f6889a50d0fdd35da2197a1c8c2a7543a600b2d6/xarray/tests/test_dataset.py#L3513): [F821](https://duckduckgo.com/?q=pep8%20F821) undefined name 'data1'
> [Line 3513:33](https://github.com/pydata/xarray/blob/f6889a50d0fdd35da2197a1c8c2a7543a600b2d6/xarray/tests/test_dataset.py#L3513): [F821](https://duckduckgo.com/?q=pep8%20F821) undefined name 'data2'
> [Line 3516:13](https://github.com/pydata/xarray/blob/f6889a50d0fdd35da2197a1c8c2a7543a600b2d6/xarray/tests/test_dataset.py#L3516): [F821](https://duckduckgo.com/?q=pep8%20F821) undefined name 'data1'


keewis: github had trouble with people abusing the "run actions on PRs" feature (namely, it was used to mine cryptocurrencies) so they now require a maintainer to approve every new commit on a first time contributor's PR (for maintainers there's a "Approve and run" button directly above the merge button)
mathause: thanks for the info, hadn't seen that one.
gcaria: Thanks all for the feedback !
This thread it's getting a bit hard to follow, but I believe this is the only suggested change that I haven't applied https://github.com/pydata/xarray/pull/5216#discussion_r624062989

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index bf0a77be..2cffe076 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -23,6 +23,8 @@ v0.17.1 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Allow passing multiple arrays to :py:meth:`Dataset.__setitem__` (:pull:`5216`).
+  By `Giacomo Caria <https://github.com/gcaria>`_.
 - Add 'cumulative' option to :py:meth:`Dataset.integrate` and
   :py:meth:`DataArray.integrate` so that result is a cumulative integral, like
   :py:meth:`scipy.integrate.cumulative_trapezoidal` (:pull:`5153`).
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index dd0a47ee..bab94f1b 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -1475,14 +1475,24 @@ class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
         else:
             return self._copy_listed(np.asarray(key))
 
+    @overload
+    def __setitem__(self, key: List[Hashable], value) -> None:
+        ...
+
+    @overload
     def __setitem__(self, key: Hashable, value) -> None:
+        ...
+
+    def __setitem__(self, key, value) -> None:
         """Add an array to this dataset.
+        Multiple arrays can be added at the same time, in which case each of
+        the following operations is applied to the respective value.
 
         If value is a `DataArray`, call its `select_vars()` method, rename it
         to `key` and merge the contents of the resulting dataset into this
         dataset.
 
-        If value is an `Variable` object (or tuple of form
+        If value is a `Variable` object (or tuple of form
         ``(dims, data[, attrs])``), add it to this dataset as a new
         variable.
         """
@@ -1491,7 +1501,27 @@ class Dataset(DataWithCoords, DatasetArithmetic, Mapping):
                 "cannot yet use a dictionary as a key to set Dataset values"
             )
 
-        self.update({key: value})
+        if isinstance(key, list):
+            if len(key) == 0:
+                raise ValueError("Empty list of variables to be set")
+            if len(key) == 1:
+                self.update({key[0]: value})
+            else:
+                if len(key) != len(value):
+                    raise ValueError(
+                        f"Different lengths of variables to be set "
+                        f"({len(key)}) and data used as input for "
+                        f"setting ({len(value)})"
+                    )
+                if isinstance(value, Dataset):
+                    self.update(dict(zip(key, value.data_vars.values())))
+                elif isinstance(value, xr.DataArray):
+                    raise ValueError("Cannot assign single DataArray to multiple keys")
+                else:
+                    self.update(dict(zip(key, value)))
+
+        else:
+            self.update({key: value})
 
     def __delitem__(self, key: Hashable) -> None:
         """Remove a variable from this dataset."""


Test patch:
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index af44e5e8..b1269369 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -3495,6 +3495,40 @@ class TestDataset:
 
         assert np.issubdtype(ds.x.dtype, dtype)
 
+    def test_setitem_using_list(self):
+
+        # assign a list of variables
+        var1 = Variable(["dim1"], np.random.randn(8))
+        var2 = Variable(["dim1"], np.random.randn(8))
+        actual = create_test_data()
+        expected = actual.copy()
+        expected["A"] = var1
+        expected["B"] = var2
+        actual[["A", "B"]] = [var1, var2]
+        assert_identical(actual, expected)
+        # assign a list of dataset arrays
+        dv = 2 * expected[["A", "B"]]
+        actual[["C", "D"]] = [d.variable for d in dv.data_vars.values()]
+        expected[["C", "D"]] = dv
+        assert_identical(actual, expected)
+
+    @pytest.mark.parametrize(
+        "var_list, data, error_regex",
+        [
+            (
+                ["A", "B"],
+                [Variable(["dim1"], np.random.randn(8))],
+                r"Different lengths",
+            ),
+            ([], [Variable(["dim1"], np.random.randn(8))], r"Empty list of variables"),
+            (["A", "B"], xr.DataArray([1, 2]), r"assign single DataArray"),
+        ],
+    )
+    def test_setitem_using_list_errors(self, var_list, data, error_regex):
+        actual = create_test_data()
+        with pytest.raises(ValueError, match=error_regex):
+            actual[var_list] = data
+
     def test_assign(self):
         ds = Dataset()
         actual = ds.assign(x=[0, 1, 2], y=2)