diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index ae42718fef3..c0f214be4c5 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -112,6 +112,37 @@ def test_getitem_with_mask(self): assert_identical(v._getitem_with_mask([0, -1, 1], fill_value=-99), self.cls(['x'], [0, -99, 1])) + def test_getitem_with_mask_fancy(self): + def one_test(key): + # key should be a dict + v = self.cls(['x', 'y', 'z'], np.random.randn(4, 5, 6)) + actual = v._getitem_with_mask(key) + + valid_indices = (~np.isnan(actual.values)).nonzero() + for indices in zip(*valid_indices): + # all the not-nan key should be nonnegative. + assert all(key[d][index] >= 0 for d, index in + zip(actual.dims, indices) if d in key.keys()) + actual[indices] = np.nan + assert np.isnan(actual.values).all() + + # orthogonal indexing + one_test({'x': [0, 1], 'y': [0, 1, 2], 'z': [3]}) + one_test({'x': [0, 1]}) + one_test({'x': [0, 1, -1], 'y': [0, 1, 2], 'z': [3]}) + # vectorized indexing + one_test({'x': Variable('a', [0, 1, -1]), + 'y': Variable('a', [0, 1, 2])}) + one_test({'x': Variable('a', [0, 1, -1]), + 'y': [-1, 2], + 'z': Variable('a', [0, 1, 2])}) + # advanced indexing + one_test({'x': Variable(('a', 'b'), [[0, 1, -1], [3, 1, 2]]), + 'y': Variable('b', [0, -1, 2])}) + one_test({'x': Variable(('a', 'b'), [[0, 1, -1], [3, 1, 2]]), + 'y': Variable('b', [0, -1, 2]), + 'z': Variable('a', [2, 1])}) + def test_getitem_with_mask_size_zero(self): v = self.cls(['x'], []) assert_identical(v._getitem_with_mask(-1), Variable((), np.nan)) @@ -1615,6 +1646,10 @@ def test_getitem_fancy(self): def test_getitem_uint(self): super(TestIndexVariable, self).test_getitem_fancy() + @pytest.mark.xfail + def test_getitem_with_mask_fancy(self): + super(TestIndexVariable, self).test_getitem_with_mask_fancy() + class TestAsCompatibleData(TestCase): def test_unchanged_types(self):