Enable __setitem__ for dask versions that support it.

xarray/core/indexing.py

@@ -4,11 +4,17 @@
 from collections import defaultdict
 from contextlib import suppress
 from datetime import timedelta
+from distutils.version import LooseVersion
 from typing import Any, Callable, Iterable, List, Sequence, Tuple, Union
 
 import numpy as np
 import pandas as pd
 
+try:
+    from dask import __version__ as dask_version
+except ModuleNotFoundError:
+    pass
+
 from . import duck_array_ops, nputils, utils
 from .npcompat import DTypeLike
 from .pycompat import (
@@ -1380,13 +1386,21 @@ def __getitem__(self, key):
                 return value
 
     def __setitem__(self, key, value):
-        raise TypeError(
-            "this variable's data is stored in a dask array, "
-            "which does not support item assignment. To "
-            "assign to this variable, you must first load it "
-            "into memory explicitly using the .load() "
-            "method or accessing its .values attribute."
-        )
+        if LooseVersion(dask_version) >= LooseVersion("2021.04.0+17"):
+            if isinstance(key, BasicIndexer):
+                self.array[key.tuple] = value
+            elif isinstance(key, VectorizedIndexer):
+                self.array.vindex[key.tuple] = value
+            elif isinstance(key, OuterIndexer):
+                self.array[key.tuple] = value
+        else:
+            raise TypeError(
+                "This variable's data is stored in a dask array, "
+                "and the installed dask version does not support item "
+                "assignment. To assign to this variable, you must either upgrade dask or"
+                "first load the variable into memory explicitly using the .load() "
+                "method or accessing its .values attribute."
+            )
 
     def transpose(self, order):
         return self.array.transpose(order)

xarray/tests/test_dask.py

@@ -111,8 +111,38 @@ def test_indexing(self):
         self.assertLazyAndIdentical(u[0], v[0])
         self.assertLazyAndIdentical(u[:1], v[:1])
         self.assertLazyAndIdentical(u[[0, 1], [0, 1, 2]], v[[0, 1], [0, 1, 2]])
-        with raises_regex(TypeError, "stored in a dask array"):
-            v[:1] = 0
+        if LooseVersion(dask.__version__) >= LooseVersion("2021.04.0+17"):
+            arr = Variable(("x"), da.array([1, 2, 3, 4]))
+            expected = Variable(("x"), da.array([99, 2, 3, 4]))
+            arr[0] = 99
+            assert_identical(arr, expected)
+            arr = Variable(("x"), da.array([1, 2, 3, 4]))
+            expected = Variable(("x"), da.array([99, 99, 99, 4]))
+            arr[2::-1] = 99
+            assert_identical(arr, expected)
+            arr = Variable(("x"), da.array([1, 2, 3, 4]))
+            expected = Variable(("x"), da.array([99, 99, 3, 99]))
+            arr[[0, -1, 1]] = 99
+            assert_identical(arr, expected)
+            arr = Variable(("x"), da.array([1, 2, 3, 4]))
+            expected = Variable(("x"), da.array([99, 99, 99, 4]))
+            arr[np.arange(3)] = 99
+            assert_identical(arr, expected)
+            arr = Variable(("x"), da.array([1, 2, 3, 4]))
+            expected = Variable(("x"), da.array([1, 99, 99, 99]))
+            arr[[False, True, True, True]] = 99
+            assert_identical(arr, expected)
+            arr = Variable(("x"), da.array([1, 2, 3, 4]))
+            expected = Variable(("x"), da.array([1, 99, 99, 99]))
+            arr[np.arange(4) > 0] = 99
+            assert_identical(arr, expected)
+            arr = Variable(("x"), da.array([1, 2, 3, 4]))
+            expected = Variable(("x"), da.array([99, 99, 99, 99]))
+            arr[arr > 0] = 99
+            assert_identical(arr, expected)
+        else:
+            with raises_regex(TypeError, "stored in a dask array"):
+                v[:1] = 0
 
     def test_squeeze(self):
         u = self.eager_var