import numpy as np from numpy.testing import assert_array_equal from xarray.core.nputils import NumpyVIndexAdapter, _is_contiguous def test_is_contiguous() -> None: assert _is_contiguous([1]) assert _is_contiguous([1, 2, 3]) assert not _is_contiguous([1, 3]) def test_vindex() -> None: x = np.arange(3 * 4 * 5).reshape((3, 4, 5)) vindex = NumpyVIndexAdapter(x) # getitem assert_array_equal(vindex[0], x[0]) assert_array_equal(vindex[[1, 2], [1, 2]], x[[1, 2], [1, 2]]) assert vindex[[0, 1], [0, 1], :].shape == (2, 5) assert vindex[[0, 1], :, [0, 1]].shape == (2, 4) assert vindex[:, [0, 1], [0, 1]].shape == (2, 3) # setitem vindex[:] = 0 assert_array_equal(x, np.zeros_like(x)) # assignment should not raise vindex[[0, 1], [0, 1], :] = vindex[[0, 1], [0, 1], :] vindex[[0, 1], :, [0, 1]] = vindex[[0, 1], :, [0, 1]] vindex[:, [0, 1], [0, 1]] = vindex[:, [0, 1], [0, 1]]