from __future__ import annotations

import pickle

import re

import sys

import warnings

from collections.abc import Hashable

from copy import deepcopy

from textwrap import dedent

from typing import Any, Final, Literal, cast

import numpy as np

import pandas as pd

import pytest

# remove once numpy 2.0 is the oldest supported version

try:

from numpy.exceptions import RankWarning

except ImportError:

from numpy import RankWarning # type: ignore[no-redef,attr-defined,unused-ignore]

import xarray as xr

from xarray import (

DataArray,

Dataset,

IndexVariable,

Variable,

align,

broadcast,

set_options,

)

from xarray.coding.times import CFDatetimeCoder

from xarray.core import dtypes

from xarray.core.common import full_like

from xarray.core.coordinates import Coordinates

from xarray.core.indexes import Index, PandasIndex, filter_indexes_from_coords

from xarray.core.types import QueryEngineOptions, QueryParserOptions

from xarray.core.utils import is_scalar

from xarray.testing import _assert_internal_invariants

from xarray.tests import (

InaccessibleArray,

ReturnItem,

assert_allclose,

assert_array_equal,

assert_chunks_equal,

assert_equal,

assert_identical,

assert_no_warnings,

has_dask,

raise_if_dask_computes,

requires_bottleneck,

requires_cupy,

requires_dask,

requires_dask_expr,

requires_iris,

requires_numexpr,

requires_pint,

requires_scipy,

requires_sparse,

source_ndarray,

)

try:

from pandas.errors import UndefinedVariableError

except ImportError:

# TODO: remove once we stop supporting pandas<1.4.3

from pandas.core.computation.ops import UndefinedVariableError

pytestmark = [

pytest.mark.filterwarnings("error:Mean of empty slice"),

pytest.mark.filterwarnings("error:All-NaN (slice|axis) encountered"),

]

class TestDataArray:

@pytest.fixture(autouse=True)

def setup(self):

self.attrs = {"attr1": "value1", "attr2": 2929}

self.x = np.random.random((10, 20))

self.v = Variable(["x", "y"], self.x)

self.va = Variable(["x", "y"], self.x, self.attrs)

self.ds = Dataset({"foo": self.v})

self.dv = self.ds["foo"]

self.mindex = pd.MultiIndex.from_product(

[["a", "b"], [1, 2]], names=("level_1", "level_2")

)

self.mda = DataArray([0, 1, 2, 3], coords={"x": self.mindex}, dims="x").astype(

np.uint64

)

def test_repr(self) -> None:

v = Variable(["time", "x"], [[1, 2, 3], [4, 5, 6]], {"foo": "bar"})

v = v.astype(np.uint64)

coords = {"x": np.arange(3, dtype=np.uint64), "other": np.uint64(0)}

data_array = DataArray(v, coords, name="my_variable")

expected = dedent(

"""\

<xarray.DataArray 'my_variable' (time: 2, x: 3)> Size: 48B

array([[1, 2, 3],

[4, 5, 6]], dtype=uint64)

Coordinates:

* x (x) uint64 24B 0 1 2

other uint64 8B 0

Dimensions without coordinates: time

Attributes:

foo: bar"""

)

assert expected == repr(data_array)

def test_repr_multiindex(self) -> None:

obj_size = np.dtype("O").itemsize

expected = dedent(

f"""\

<xarray.DataArray (x: 4)> Size: 32B

array([0, 1, 2, 3], dtype=uint64)

Coordinates:

* x (x) object {4 * obj_size}B MultiIndex

* level_1 (x) object {4 * obj_size}B 'a' 'a' 'b' 'b'

* level_2 (x) int64 32B 1 2 1 2"""

)

assert expected == repr(self.mda)

def test_repr_multiindex_long(self) -> None:

mindex_long = pd.MultiIndex.from_product(

[["a", "b", "c", "d"], [1, 2, 3, 4, 5, 6, 7, 8]],

names=("level_1", "level_2"),

)

mda_long = DataArray(

list(range(32)), coords={"x": mindex_long}, dims="x"

).astype(np.uint64)

obj_size = np.dtype("O").itemsize

expected = dedent(

f"""\

<xarray.DataArray (x: 32)> Size: 256B

array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,

17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31],

dtype=uint64)

Coordinates:

* x (x) object {32 * obj_size}B MultiIndex

* level_1 (x) object {32 * obj_size}B 'a' 'a' 'a' 'a' 'a' 'a' ... 'd' 'd' 'd' 'd' 'd' 'd'

* level_2 (x) int64 256B 1 2 3 4 5 6 7 8 1 2 3 4 ... 5 6 7 8 1 2 3 4 5 6 7 8"""

)

assert expected == repr(mda_long)

def test_properties(self) -> None:

assert_equal(self.dv.variable, self.v)

assert_array_equal(self.dv.values, self.v.values)

for attr in ["dims", "dtype", "shape", "size", "nbytes", "ndim", "attrs"]:

assert getattr(self.dv, attr) == getattr(self.v, attr)

assert len(self.dv) == len(self.v)

assert_equal(self.dv.variable, self.v)

assert set(self.dv.coords) == set(self.ds.coords)

for k, v in self.dv.coords.items():

assert_array_equal(v, self.ds.coords[k])

with pytest.raises(AttributeError):

self.dv.dataset

assert isinstance(self.ds["x"].to_index(), pd.Index)

with pytest.raises(ValueError, match=r"must be 1-dimensional"):

self.ds["foo"].to_index()

with pytest.raises(AttributeError):

self.dv.variable = self.v

def test_data_property(self) -> None:

array = DataArray(np.zeros((3, 4)))

actual = array.copy()

actual.values = np.ones((3, 4))

assert_array_equal(np.ones((3, 4)), actual.values)

actual.data = 2 * np.ones((3, 4))

assert_array_equal(2 * np.ones((3, 4)), actual.data)

assert_array_equal(actual.data, actual.values)

def test_indexes(self) -> None:

array = DataArray(np.zeros((2, 3)), [("x", [0, 1]), ("y", ["a", "b", "c"])])

expected_indexes = {"x": pd.Index([0, 1]), "y": pd.Index(["a", "b", "c"])}

expected_xindexes = {

k: PandasIndex(idx, k) for k, idx in expected_indexes.items()

}

assert array.xindexes.keys() == expected_xindexes.keys()

assert array.indexes.keys() == expected_indexes.keys()

assert all([isinstance(idx, pd.Index) for idx in array.indexes.values()])

assert all([isinstance(idx, Index) for idx in array.xindexes.values()])

for k in expected_indexes:

assert array.xindexes[k].equals(expected_xindexes[k])

assert array.indexes[k].equals(expected_indexes[k])

def test_get_index(self) -> None:

array = DataArray(np.zeros((2, 3)), coords={"x": ["a", "b"]}, dims=["x", "y"])

assert array.get_index("x").equals(pd.Index(["a", "b"]))

assert array.get_index("y").equals(pd.Index([0, 1, 2]))

with pytest.raises(KeyError):

array.get_index("z")

def test_get_index_size_zero(self) -> None:

array = DataArray(np.zeros((0,)), dims=["x"])

actual = array.get_index("x")

expected = pd.Index([], dtype=np.int64)

assert actual.equals(expected)

assert actual.dtype == expected.dtype

def test_struct_array_dims(self) -> None:

"""

This test checks subtraction of two DataArrays for the case

when dimension is a structured array.

"""

# GH837, GH861

# checking array subtraction when dims are the same

p_data = np.array(

[("Abe", 180), ("Stacy", 150), ("Dick", 200)],

dtype=[("name", "|S256"), ("height", object)],

)

weights_0 = DataArray(

[80, 56, 120], dims=["participant"], coords={"participant": p_data}

)

weights_1 = DataArray(

[81, 52, 115], dims=["participant"], coords={"participant": p_data}

)

actual = weights_1 - weights_0

expected = DataArray(

[1, -4, -5], dims=["participant"], coords={"participant": p_data}

)

assert_identical(actual, expected)

# checking array subtraction when dims are not the same

p_data_alt = np.array(

[("Abe", 180), ("Stacy", 151), ("Dick", 200)],

dtype=[("name", "|S256"), ("height", object)],

)

weights_1 = DataArray(

[81, 52, 115], dims=["participant"], coords={"participant": p_data_alt}

)

actual = weights_1 - weights_0

expected = DataArray(

[1, -5], dims=["participant"], coords={"participant": p_data[[0, 2]]}

)

assert_identical(actual, expected)

# checking array subtraction when dims are not the same and one

# is np.nan

p_data_nan = np.array(

[("Abe", 180), ("Stacy", np.nan), ("Dick", 200)],

dtype=[("name", "|S256"), ("height", object)],

)

weights_1 = DataArray(

[81, 52, 115], dims=["participant"], coords={"participant": p_data_nan}

)

actual = weights_1 - weights_0

expected = DataArray(

[1, -5], dims=["participant"], coords={"participant": p_data[[0, 2]]}

)

assert_identical(actual, expected)

def test_name(self) -> None:

arr = self.dv

assert arr.name == "foo"

copied = arr.copy()

arr.name = "bar"

assert arr.name == "bar"

assert_equal(copied, arr)

actual = DataArray(IndexVariable("x", [3]))

actual.name = "y"

expected = DataArray([3], [("x", [3])], name="y")

assert_identical(actual, expected)

def test_dims(self) -> None:

arr = self.dv

assert arr.dims == ("x", "y")

with pytest.raises(AttributeError, match=r"you cannot assign"):

arr.dims = ("w", "z")

def test_sizes(self) -> None:

array = DataArray(np.zeros((3, 4)), dims=["x", "y"])

assert array.sizes == {"x": 3, "y": 4}

assert tuple(array.sizes) == array.dims

with pytest.raises(TypeError):

array.sizes["foo"] = 5 # type: ignore

def test_encoding(self) -> None:

expected = {"foo": "bar"}

self.dv.encoding["foo"] = "bar"

assert expected == self.dv.encoding

expected2 = {"baz": 0}

self.dv.encoding = expected2

assert expected2 is not self.dv.encoding

def test_drop_encoding(self) -> None:

array = self.mda

encoding = {"scale_factor": 10}

array.encoding = encoding

array["x"].encoding = encoding

assert array.encoding == encoding

assert array["x"].encoding == encoding

actual = array.drop_encoding()

# did not modify in place

assert array.encoding == encoding

assert array["x"].encoding == encoding

# variable and coord encoding is empty

assert actual.encoding == {}

assert actual["x"].encoding == {}

def test_constructor(self) -> None:

data = np.random.random((2, 3))

# w/o coords, w/o dims

actual = DataArray(data)

expected = Dataset({None: (["dim_0", "dim_1"], data)})[None]

assert_identical(expected, actual)

actual = DataArray(data, [["a", "b"], [-1, -2, -3]])

expected = Dataset(

{

None: (["dim_0", "dim_1"], data),

"dim_0": ("dim_0", ["a", "b"]),

"dim_1": ("dim_1", [-1, -2, -3]),

}

)[None]

assert_identical(expected, actual)

# pd.Index coords, w/o dims

actual = DataArray(

data, [pd.Index(["a", "b"], name="x"), pd.Index([-1, -2, -3], name="y")]

)

expected = Dataset(

{None: (["x", "y"], data), "x": ("x", ["a", "b"]), "y": ("y", [-1, -2, -3])}

)[None]

assert_identical(expected, actual)

# list coords, w dims

coords1: list[Any] = [["a", "b"], [-1, -2, -3]]

actual = DataArray(data, coords1, ["x", "y"])

assert_identical(expected, actual)

# pd.Index coords, w dims

coords2: list[pd.Index] = [

pd.Index(["a", "b"], name="A"),

pd.Index([-1, -2, -3], name="B"),

]

actual = DataArray(data, coords2, ["x", "y"])

assert_identical(expected, actual)

# dict coords, w dims

coords3 = {"x": ["a", "b"], "y": [-1, -2, -3]}

actual = DataArray(data, coords3, ["x", "y"])

assert_identical(expected, actual)

# dict coords, w/o dims

actual = DataArray(data, coords3)

assert_identical(expected, actual)

# tuple[dim, list] coords, w/o dims

coords4 = [("x", ["a", "b"]), ("y", [-1, -2, -3])]

actual = DataArray(data, coords4)

assert_identical(expected, actual)

# partial dict coords, w dims

expected = Dataset({None: (["x", "y"], data), "x": ("x", ["a", "b"])})[None]

actual = DataArray(data, {"x": ["a", "b"]}, ["x", "y"])

assert_identical(expected, actual)

# w/o coords, w dims

actual = DataArray(data, dims=["x", "y"])

expected = Dataset({None: (["x", "y"], data)})[None]

assert_identical(expected, actual)

# w/o coords, w dims, w name

actual = DataArray(data, dims=["x", "y"], name="foo")

expected = Dataset({"foo": (["x", "y"], data)})["foo"]

assert_identical(expected, actual)

# w/o coords, w/o dims, w name

actual = DataArray(data, name="foo")

expected = Dataset({"foo": (["dim_0", "dim_1"], data)})["foo"]

assert_identical(expected, actual)

# w/o coords, w dims, w attrs

actual = DataArray(data, dims=["x", "y"], attrs={"bar": 2})

expected = Dataset({None: (["x", "y"], data, {"bar": 2})})[None]

assert_identical(expected, actual)

# w/o coords, w dims (ds has attrs)

actual = DataArray(data, dims=["x", "y"])

expected = Dataset({None: (["x", "y"], data, {}, {"bar": 2})})[None]

assert_identical(expected, actual)

# data is list, w coords

actual = DataArray([1, 2, 3], coords={"x": [0, 1, 2]})

expected = DataArray([1, 2, 3], coords=[("x", [0, 1, 2])])

assert_identical(expected, actual)

def test_constructor_invalid(self) -> None:

data = np.random.randn(3, 2)

with pytest.raises(ValueError, match=r"coords is not dict-like"):

DataArray(data, [[0, 1, 2]], ["x", "y"])

with pytest.raises(ValueError, match=r"not a subset of the .* dim"):

DataArray(data, {"x": [0, 1, 2]}, ["a", "b"])

with pytest.raises(ValueError, match=r"not a subset of the .* dim"):

DataArray(data, {"x": [0, 1, 2]})

with pytest.raises(TypeError, match=r"is not hashable"):

DataArray(data, dims=["x", []]) # type: ignore[list-item]

with pytest.raises(ValueError, match=r"conflicting sizes for dim"):

DataArray([1, 2, 3], coords=[("x", [0, 1])])

with pytest.raises(ValueError, match=r"conflicting sizes for dim"):

DataArray([1, 2], coords={"x": [0, 1], "y": ("x", [1])}, dims="x")

with pytest.raises(ValueError, match=r"conflicting MultiIndex"):

DataArray(np.random.rand(4, 4), [("x", self.mindex), ("y", self.mindex)])

with pytest.raises(ValueError, match=r"conflicting MultiIndex"):

DataArray(np.random.rand(4, 4), [("x", self.mindex), ("level_1", range(4))])

def test_constructor_from_self_described(self) -> None:

data: list[list[float]] = [[-0.1, 21], [0, 2]]

expected = DataArray(

data,

coords={"x": ["a", "b"], "y": [-1, -2]},

dims=["x", "y"],

name="foobar",

attrs={"bar": 2},

)

actual = DataArray(expected)

assert_identical(expected, actual)

actual = DataArray(expected.values, actual.coords)

assert_equal(expected, actual)

frame = pd.DataFrame(

data,

index=pd.Index(["a", "b"], name="x"),

columns=pd.Index([-1, -2], name="y"),

)

actual = DataArray(frame)

assert_equal(expected, actual)

series = pd.Series(data[0], index=pd.Index([-1, -2], name="y"))

actual = DataArray(series)

assert_equal(expected[0].reset_coords("x", drop=True), actual)

expected = DataArray(

data,

coords={"x": ["a", "b"], "y": [-1, -2], "a": 0, "z": ("x", [-0.5, 0.5])},

dims=["x", "y"],

)

actual = DataArray(expected)

assert_identical(expected, actual)

actual = DataArray(expected.values, expected.coords)

assert_identical(expected, actual)

expected = Dataset({"foo": ("foo", ["a", "b"])})["foo"]

actual = DataArray(pd.Index(["a", "b"], name="foo"))

assert_identical(expected, actual)

actual = DataArray(IndexVariable("foo", ["a", "b"]))

assert_identical(expected, actual)

@requires_dask

def test_constructor_from_self_described_chunked(self) -> None:

expected = DataArray(

[[-0.1, 21], [0, 2]],

coords={"x": ["a", "b"], "y": [-1, -2]},

dims=["x", "y"],

name="foobar",

attrs={"bar": 2},

).chunk()

actual = DataArray(expected)

assert_identical(expected, actual)

assert_chunks_equal(expected, actual)

def test_constructor_from_0d(self) -> None:

expected = Dataset({None: ([], 0)})[None]

actual = DataArray(0)

assert_identical(expected, actual)

@requires_dask

def test_constructor_dask_coords(self) -> None:

# regression test for GH1684

import dask.array as da

coord = da.arange(8, chunks=(4,))

data = da.random.random((8, 8), chunks=(4, 4)) + 1

actual = DataArray(data, coords={"x": coord, "y": coord}, dims=["x", "y"])

ecoord = np.arange(8)

expected = DataArray(data, coords={"x": ecoord, "y": ecoord}, dims=["x", "y"])

assert_equal(actual, expected)

def test_constructor_no_default_index(self) -> None:

# explicitly passing a Coordinates object skips the creation of default index

da = DataArray(range(3), coords=Coordinates({"x": [1, 2, 3]}, indexes={}))

assert "x" in da.coords

assert "x" not in da.xindexes

def test_constructor_multiindex(self) -> None:

midx = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=("one", "two"))

coords = Coordinates.from_pandas_multiindex(midx, "x")

da = DataArray(range(4), coords=coords, dims="x")

assert_identical(da.coords, coords)

def test_constructor_custom_index(self) -> None:

class CustomIndex(Index): ...

coords = Coordinates(

coords={"x": ("x", [1, 2, 3])}, indexes={"x": CustomIndex()}

)

da = DataArray(range(3), coords=coords)

assert isinstance(da.xindexes["x"], CustomIndex)

# test coordinate variables copied

assert da.coords["x"] is not coords.variables["x"]

def test_equals_and_identical(self) -> None:

orig = DataArray(np.arange(5.0), {"a": 42}, dims="x")

expected = orig

actual = orig.copy()

assert expected.equals(actual)

assert expected.identical(actual)

actual = expected.rename("baz")

assert expected.equals(actual)

assert not expected.identical(actual)

actual = expected.rename({"x": "xxx"})

assert not expected.equals(actual)

assert not expected.identical(actual)

actual = expected.copy()

actual.attrs["foo"] = "bar"

assert expected.equals(actual)

assert not expected.identical(actual)

actual = expected.copy()

actual["x"] = ("x", -np.arange(5))

assert not expected.equals(actual)

assert not expected.identical(actual)

actual = expected.reset_coords(drop=True)

assert not expected.equals(actual)

assert not expected.identical(actual)

actual = orig.copy()

actual[0] = np.nan

expected = actual.copy()

assert expected.equals(actual)

assert expected.identical(actual)

actual[:] = np.nan

assert not expected.equals(actual)

assert not expected.identical(actual)

actual = expected.copy()

actual["a"] = 100000

assert not expected.equals(actual)

assert not expected.identical(actual)

def test_equals_failures(self) -> None:

orig = DataArray(np.arange(5.0), {"a": 42}, dims="x")

assert not orig.equals(np.arange(5)) # type: ignore[arg-type]

assert not orig.identical(123) # type: ignore

assert not orig.broadcast_equals({1: 2}) # type: ignore

def test_broadcast_equals(self) -> None:

a = DataArray([0, 0], {"y": 0}, dims="x")

b = DataArray([0, 0], {"y": ("x", [0, 0])}, dims="x")

assert a.broadcast_equals(b)

assert b.broadcast_equals(a)

assert not a.equals(b)

assert not a.identical(b)

c = DataArray([0], coords={"x": 0}, dims="y")

assert not a.broadcast_equals(c)

assert not c.broadcast_equals(a)

def test_getitem(self) -> None:

# strings pull out dataarrays

assert_identical(self.dv, self.ds["foo"])

x = self.dv["x"]

y = self.dv["y"]

assert_identical(self.ds["x"], x)

assert_identical(self.ds["y"], y)

arr = ReturnItem()

for i in [

arr[:],

arr[...],

arr[x.values],

arr[x.variable],

arr[x],

arr[x, y],

arr[x.values > -1],

arr[x.variable > -1],

arr[x > -1],

arr[x > -1, y > -1],

]:

assert_equal(self.dv, self.dv[i])

for i in [

arr[0],

arr[:, 0],

arr[:3, :2],

arr[x.values[:3]],

arr[x.variable[:3]],

arr[x[:3]],

arr[x[:3], y[:4]],

arr[x.values > 3],

arr[x.variable > 3],

arr[x > 3],

arr[x > 3, y > 3],

]:

assert_array_equal(self.v[i], self.dv[i])

def test_getitem_dict(self) -> None:

actual = self.dv[{"x": slice(3), "y": 0}]

expected = self.dv.isel(x=slice(3), y=0)

assert_identical(expected, actual)

def test_getitem_coords(self) -> None:

orig = DataArray(

[[10], [20]],

{

"x": [1, 2],

"y": [3],

"z": 4,

"x2": ("x", ["a", "b"]),

"y2": ("y", ["c"]),

"xy": (["y", "x"], [["d", "e"]]),

},

dims=["x", "y"],

)

assert_identical(orig, orig[:])

assert_identical(orig, orig[:, :])

assert_identical(orig, orig[...])

assert_identical(orig, orig[:2, :1])

assert_identical(orig, orig[[0, 1], [0]])

actual = orig[0, 0]

expected = DataArray(

10, {"x": 1, "y": 3, "z": 4, "x2": "a", "y2": "c", "xy": "d"}

)

assert_identical(expected, actual)

actual = orig[0, :]

expected = DataArray(

[10],

{

"x": 1,

"y": [3],

"z": 4,

"x2": "a",

"y2": ("y", ["c"]),

"xy": ("y", ["d"]),

},

dims="y",

)

assert_identical(expected, actual)

actual = orig[:, 0]

expected = DataArray(

[10, 20],

{

"x": [1, 2],

"y": 3,

"z": 4,

"x2": ("x", ["a", "b"]),

"y2": "c",

"xy": ("x", ["d", "e"]),

},

dims="x",

)

assert_identical(expected, actual)

def test_getitem_dataarray(self) -> None:

# It should not conflict

da = DataArray(np.arange(12).reshape((3, 4)), dims=["x", "y"])

ind = DataArray([[0, 1], [0, 1]], dims=["x", "z"])

actual = da[ind]

assert_array_equal(actual, da.values[[[0, 1], [0, 1]], :])

da = DataArray(

np.arange(12).reshape((3, 4)),

dims=["x", "y"],

coords={"x": [0, 1, 2], "y": ["a", "b", "c", "d"]},

)

ind = xr.DataArray([[0, 1], [0, 1]], dims=["X", "Y"])

actual = da[ind]

expected = da.values[[[0, 1], [0, 1]], :]

assert_array_equal(actual, expected)

assert actual.dims == ("X", "Y", "y")

# boolean indexing

ind = xr.DataArray([True, True, False], dims=["x"])

assert_equal(da[ind], da[[0, 1], :])

assert_equal(da[ind], da[[0, 1]])

assert_equal(da[ind], da[ind.values])

def test_getitem_empty_index(self) -> None:

da = DataArray(np.arange(12).reshape((3, 4)), dims=["x", "y"])

assert_identical(da[{"x": []}], DataArray(np.zeros((0, 4)), dims=["x", "y"]))

assert_identical(

da.loc[{"y": []}], DataArray(np.zeros((3, 0)), dims=["x", "y"])

)

assert_identical(da[[]], DataArray(np.zeros((0, 4)), dims=["x", "y"]))

def test_setitem(self) -> None:

# basic indexing should work as numpy's indexing

tuples = [

(0, 0),

(0, slice(None, None)),

(slice(None, None), slice(None, None)),

(slice(None, None), 0),

([1, 0], slice(None, None)),

(slice(None, None), [1, 0]),

]

for t in tuples:

expected = np.arange(6).reshape(3, 2)

orig = DataArray(

np.arange(6).reshape(3, 2),

{

"x": [1, 2, 3],

"y": ["a", "b"],

"z": 4,

"x2": ("x", ["a", "b", "c"]),

"y2": ("y", ["d", "e"]),

},

dims=["x", "y"],

)

orig[t] = 1

expected[t] = 1

assert_array_equal(orig.values, expected)

def test_setitem_fancy(self) -> None:

# vectorized indexing

da = DataArray(np.ones((3, 2)), dims=["x", "y"])

ind = Variable(["a"], [0, 1])

da[dict(x=ind, y=ind)] = 0

expected = DataArray([[0, 1], [1, 0], [1, 1]], dims=["x", "y"])

assert_identical(expected, da)

# assign another 0d-variable

da[dict(x=ind, y=ind)] = Variable((), 0)

expected = DataArray([[0, 1], [1, 0], [1, 1]], dims=["x", "y"])

assert_identical(expected, da)

# assign another 1d-variable

da[dict(x=ind, y=ind)] = Variable(["a"], [2, 3])

expected = DataArray([[2, 1], [1, 3], [1, 1]], dims=["x", "y"])

assert_identical(expected, da)

# 2d-vectorized indexing

da = DataArray(np.ones((3, 2)), dims=["x", "y"])

ind_x = DataArray([[0, 1]], dims=["a", "b"])

ind_y = DataArray([[1, 0]], dims=["a", "b"])

da[dict(x=ind_x, y=ind_y)] = 0

expected = DataArray([[1, 0], [0, 1], [1, 1]], dims=["x", "y"])

assert_identical(expected, da)

da = DataArray(np.ones((3, 2)), dims=["x", "y"])

ind = Variable(["a"], [0, 1])

da[ind] = 0

expected = DataArray([[0, 0], [0, 0], [1, 1]], dims=["x", "y"])

assert_identical(expected, da)

def test_setitem_dataarray(self) -> None:

def get_data():

return DataArray(

np.ones((4, 3, 2)),

dims=["x", "y", "z"],

coords={

"x": np.arange(4),

"y": ["a", "b", "c"],

"non-dim": ("x", [1, 3, 4, 2]),

},

)

da = get_data()

# indexer with inconsistent coordinates.

ind = DataArray(np.arange(1, 4), dims=["x"], coords={"x": np.random.randn(3)})

with pytest.raises(IndexError, match=r"dimension coordinate 'x'"):

da[dict(x=ind)] = 0

# indexer with consistent coordinates.

ind = DataArray(np.arange(1, 4), dims=["x"], coords={"x": np.arange(1, 4)})

da[dict(x=ind)] = 0 # should not raise

assert np.allclose(da[dict(x=ind)].values, 0)

assert_identical(da["x"], get_data()["x"])

assert_identical(da["non-dim"], get_data()["non-dim"])

da = get_data()

# conflict in the assigning values

value = xr.DataArray(

np.zeros((3, 3, 2)),

dims=["x", "y", "z"],

coords={"x": [0, 1, 2], "non-dim": ("x", [0, 2, 4])},

)

with pytest.raises(IndexError, match=r"dimension coordinate 'x'"):

da[dict(x=ind)] = value

# consistent coordinate in the assigning values

value = xr.DataArray(

np.zeros((3, 3, 2)),

dims=["x", "y", "z"],

coords={"x": [1, 2, 3], "non-dim": ("x", [0, 2, 4])},

)

da[dict(x=ind)] = value

assert np.allclose(da[dict(x=ind)].values, 0)

assert_identical(da["x"], get_data()["x"])

assert_identical(da["non-dim"], get_data()["non-dim"])

# Conflict in the non-dimension coordinate

value = xr.DataArray(

np.zeros((3, 3, 2)),

dims=["x", "y", "z"],

coords={"x": [1, 2, 3], "non-dim": ("x", [0, 2, 4])},

)

da[dict(x=ind)] = value # should not raise

# conflict in the assigning values

value = xr.DataArray(

np.zeros((3, 3, 2)),

dims=["x", "y", "z"],

coords={"x": [0, 1, 2], "non-dim": ("x", [0, 2, 4])},

)

with pytest.raises(IndexError, match=r"dimension coordinate 'x'"):

da[dict(x=ind)] = value

# consistent coordinate in the assigning values

value = xr.DataArray(

np.zeros((3, 3, 2)),

dims=["x", "y", "z"],

coords={"x": [1, 2, 3], "non-dim": ("x", [0, 2, 4])},

)

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"])

c = xr.DataArray([[0, 1], [2, 3]], dims=["u", "v"])

w = xr.DataArray([-1, -2], dims=["u"])

index = dict(b=b, c=c)

v[index] = w

assert (v[index] == w).all()

# Indexing with coordinates

v = xr.DataArray(np.r_[:120].reshape(2, 3, 4, 5), dims=["a", "b", "c", "d"])

v.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)

v.loc[index] = w

assert (v.loc[index] == w).all()

def test_contains(self) -> None:

data_array = DataArray([1, 2])

assert 1 in data_array

assert 3 not in data_array

def test_pickle(self) -> None:

data = DataArray(np.random.random((3, 3)), dims=("id", "time"))

roundtripped = pickle.loads(pickle.dumps(data))

assert_identical(data, roundtripped)

@requires_dask

def test_chunk(self) -> None:

unblocked = DataArray(np.ones((3, 4)))

assert unblocked.chunks is None

blocked = unblocked.chunk()

assert blocked.chunks == ((3,), (4,))

first_dask_name = blocked.data.name

with pytest.warns(DeprecationWarning):

blocked = unblocked.chunk(chunks=((2, 1), (2, 2))) # type: ignore

assert blocked.chunks == ((2, 1), (2, 2))

assert blocked.data.name != first_dask_name

blocked = unblocked.chunk(chunks=(3, 3))

assert blocked.chunks == ((3,), (3, 1))

assert blocked.data.name != first_dask_name

# name doesn't change when rechunking by same amount

# this fails if ReprObject doesn't have __dask_tokenize__ defined

assert unblocked.chunk(2).data.name == unblocked.chunk(2).data.name

assert blocked.load().chunks is None

# Check that kwargs are passed

import dask.array as da

blocked = unblocked.chunk(name_prefix="testname_")

assert isinstance(blocked.data, da.Array)

assert "testname_" in blocked.data.name

# test kwargs form of chunks

blocked = unblocked.chunk(dim_0=3, dim_1=3)

assert blocked.chunks == ((3,), (3, 1))

assert blocked.data.name != first_dask_name

def test_isel(self) -> None:

assert_identical(self.dv[0], self.dv.isel(x=0))

assert_identical(self.dv, self.dv.isel(x=slice(None)))

assert_identical(self.dv[:3], self.dv.isel(x=slice(3)))

assert_identical(self.dv[:3, :5], self.dv.isel(x=slice(3), y=slice(5)))

with pytest.raises(

ValueError,

match=r"Dimensions {'not_a_dim'} do not exist. Expected "

r"one or more of \('x', 'y'\)",

):

self.dv.isel(not_a_dim=0)

with pytest.warns(

UserWarning,

match=r"Dimensions {'not_a_dim'} do not exist. "

r"Expected one or more of \('x', 'y'\)",

):

self.dv.isel(not_a_dim=0, missing_dims="warn")

assert_identical(self.dv, self.dv.isel(not_a_dim=0, missing_dims="ignore"))

def test_isel_types(self) -> None:

# regression test for #1405

da = DataArray([1, 2, 3], dims="x")

# uint64

assert_identical(

da.isel(x=np.array([0], dtype="uint64")), da.isel(x=np.array([0]))

)

# uint32

assert_identical(

da.isel(x=np.array([0], dtype="uint32")), da.isel(x=np.array([0]))

)

# int64

assert_identical(

da.isel(x=np.array([0], dtype="int64")), da.isel(x=np.array([0]))

)

@pytest.mark.filterwarnings("ignore::DeprecationWarning")

def test_isel_fancy(self) -> None:

shape = (10, 7, 6)

np_array = np.random.random(shape)

da = DataArray(

np_array, dims=["time", "y", "x"], coords={"time": np.arange(0, 100, 10)}

)

y = [1, 3]

x = [3, 0]

expected = da.values[:, y, x]

actual = da.isel(y=(("test_coord",), y), x=(("test_coord",), x))

assert actual.coords["test_coord"].shape == (len(y),)

assert list(actual.coords) == ["time"]

assert actual.dims == ("time", "test_coord")

np.testing.assert_equal(actual, expected)

# a few corner cases

da.isel(

time=(("points",), [1, 2]), x=(("points",), [2, 2]), y=(("points",), [3, 4])

)

np.testing.assert_allclose(

da.isel(

time=(("p",), [1]), x=(("p",), [2]), y=(("p",), [4])

).values.squeeze(),

np_array[1, 4, 2].squeeze(),

)

da.isel(time=(("points",), [1, 2]))

y = [-1, 0]

x = [-2, 2]

expected2 = da.values[:, y, x]

actual2 = da.isel(x=(("points",), x), y=(("points",), y)).values

np.testing.assert_equal(actual2, expected2)

# test that the order of the indexers doesn't matter

assert_identical(

da.isel(y=(("points",), y), x=(("points",), x)),

da.isel(x=(("points",), x), y=(("points",), y)),

)

# make sure we're raising errors in the right places

with pytest.raises(IndexError, match=r"Dimensions of indexers mismatch"):

da.isel(y=(("points",), [1, 2]), x=(("points",), [1, 2, 3]))

# tests using index or DataArray as indexers

stations = Dataset()