Function test_get_dummies_broadcast

xarray/tests/test_accessor_str.py:3016–3041

(dtype)

def test_get_dummies_broadcast(dtype) -> None:
    values = xr.DataArray(
        ["x~x|x~x", "x", "x|x~x", "x~x"],
        dims=["X"],
    ).astype(dtype)

    sep = xr.DataArray(
        ["|", "~"],
        dims=["Y"],
    ).astype(dtype)

    expected_list = [
        [[False, False, True], [True, True, False]],
        [[True, False, False], [True, False, False]],
        [[True, False, True], [True, True, False]],
        [[False, False, True], [True, False, False]],
    ]
    expected_np = np.array(expected_list)
    expected = xr.DataArray(expected_np, dims=["X", "Y", "ZZ"])
    expected.coords["ZZ"] = np.array(["x", "x|x", "x~x"]).astype(dtype)

    res = values.str.get_dummies(dim="ZZ", sep=sep)

    assert res.dtype == expected.dtype

    assert_equal(res, expected)


def test_get_dummies_empty(dtype) -> None: