Add regression test for Dataset (pydata#7030)

dranjan · Aug 18, 2023 · f433db6 · f433db6
1 parent e152e4b
commit f433db6
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
@@ -835,6 +835,7 @@ def get_data():
         da[dict(x=ind)] = value  # should not raise
 
     def test_setitem_vectorized(self) -> None:
+        # Regression test for GH:7030
         # Positional indexing
         v = xr.DataArray(np.r_[:120].reshape(2, 3, 4, 5), dims=["a", "b", "c", "d"])
         b = xr.DataArray([[0, 0], [1, 0]], dims=["u", "v"])

diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
@@ -4160,6 +4160,29 @@ def test_setitem_align_new_indexes(self) -> None:
         )
         assert_identical(ds, expected)
 
+    def test_setitem_vectorized(self) -> None:
+        # Regression test for GH:7030
+        # Positional indexing
+        da = xr.DataArray(np.r_[:120].reshape(2, 3, 4, 5), dims=["a", "b", "c", "d"])
+        ds = xr.Dataset({'da': da})
+        b = xr.DataArray([[0, 0], [1, 0]], dims=["u", "v"])
+        c = xr.DataArray([[0, 1], [2, 3]], dims=["u", "v"])
+        w = xr.DataArray([-1, -2], dims=["u"])
+        index = dict(b=b, c=c)
+        ds[index] = xr.Dataset({'da': w})
+        assert (ds[index]['da'] == w).all()
+
+        # Indexing with coordinates
+        da = xr.DataArray(np.r_[:120].reshape(2, 3, 4, 5), dims=["a", "b", "c", "d"])
+        ds = xr.Dataset({'da': da})
+        ds.coords["b"] = [2, 4, 6]
+        b = xr.DataArray([[2, 2], [4, 2]], dims=["u", "v"])
+        c = xr.DataArray([[0, 1], [2, 3]], dims=["u", "v"])
+        w = xr.DataArray([-1, -2], dims=["u"])
+        index = dict(b=b, c=c)
+        ds.loc[index] = xr.Dataset({'da': w}, coords={'b': ds.coords['b']})
+        assert (ds.loc[index]['da'] == w).all()
+
     @pytest.mark.parametrize("dtype", [str, bytes])
     def test_setitem_str_dtype(self, dtype) -> None:
         ds = xr.Dataset(coords={"x": np.array(["x", "y"], dtype=dtype)})