33 lines
987 B
Python
33 lines
987 B
Python
from __future__ import annotations
|
|
|
|
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]]
|