import pytest
import numpy as np

import dask.array as da
from dask.array.numpy_compat import _make_sliced_dtype
from dask.array.utils import assert_eq


@pytest.fixture(
    params=[
        [("A", ("f4", (3, 2))), ("B", ("f4", 3)), ("C", ("f8", 3))],
        [("A", ("i4", (3, 2))), ("B", ("f4", 3)), ("C", ("S4", 3))],
    ]
)
def dtype(request):
    return np.dtype(request.param)


@pytest.fixture(params=[["A"], ["A", "B"], ["A", "B", "C"]])
def index(request):
    return request.param


def test_basic():
    # sanity check
    dtype = [("a", "f8"), ("b", "f8"), ("c", "f8")]
    x = np.ones((5, 3), dtype=dtype)
    dx = da.ones((5, 3), dtype=dtype, chunks=3)
    result = dx[["a", "b"]]
    expected = x[["a", "b"]]
    assert_eq(result, expected)


def test_slice_dtype(dtype, index):
    result = _make_sliced_dtype(dtype, index)
    expected = np.ones((5, len(dtype)), dtype=dtype)[index].dtype
    assert result == expected
