Function construct_dataarray

xarray/tests/test_duck_array_ops.py:373–402

(dim_num, dtype, contains_nan, dask)

def construct_dataarray(dim_num, dtype, contains_nan, dask):
    # dimnum <= 3
    rng = np.random.default_rng(0)
    shapes = [16, 8, 4][:dim_num]
    dims = ("x", "y", "z")[:dim_num]

    if np.issubdtype(dtype, np.floating):
        array = rng.random(shapes).astype(dtype)
    elif np.issubdtype(dtype, np.integer):
        array = rng.integers(0, 10, size=shapes).astype(dtype)
    elif np.issubdtype(dtype, np.bool_):
        array = rng.integers(0, 1, size=shapes).astype(dtype)
    elif dtype is str:
        array = rng.choice(["a", "b", "c", "d"], size=shapes)
    else:
        raise ValueError

    if contains_nan:
        inds = rng.choice(range(array.size), int(array.size * 0.2))
        dtype, fill_value = dtypes.maybe_promote(array.dtype)
        array = array.astype(dtype)
        array.flat[inds] = fill_value

    da = DataArray(array, dims=dims, coords={"x": np.arange(16)}, name="da")

    if dask and has_dask:
        chunks = dict.fromkeys(dims, 4)
        da = da.chunk(chunks)

    return da


def from_series_or_scalar(se):