rio/tests/test_table_indexing.py

"""
Tables support numpy-style 2D indexing. This is rather complex, hence the
tests here.
"""

import typing as t

import pytest

import rio


# Helper class for easily creating indices
class MakeIndex:
    def __getitem__(self, index) -> t.Any:
        return index


make_index = MakeIndex()


@pytest.mark.parametrize(
    "index,enable_column_names,result_should",
    [
        # Valid index without anything special
        (
            make_index[1:3, 4:7],
            False,
            (4, 1, 3, 2),
        ),
        # Selecting an entire column
        (
            make_index["a"],
            True,
            (0, 0, 1, 20),
        ),
        (
            make_index["b"],
            True,
            (1, 0, 1, 20),
        ),
        # Attempting to select a column, but no column names are available
        (
            make_index["a"],
            False,
            ValueError,
        ),
        # Selecting a non-existent column
        (
            make_index["invalid_column_name"],
            True,
            KeyError,
        ),
        # Selecting some rows in a column
        (
            make_index[1:3, "a"],
            True,
            (0, 1, 1, 2),
        ),
        # Attempting to select by name in the wrong axis
        (
            make_index["a", 0],
            True,
            ValueError,
        ),
        (
            make_index["a", "b"],
            True,
            ValueError,
        ),
        # Wrong number of indices
        (
            make_index[0],
            True,
            ValueError,
        ),
        (
            make_index[1, 2, 3],
            True,
            ValueError,
        ),
        (
            make_index[1, 2, 3, 4],
            True,
            ValueError,
        ),
        # Wrong datatypes
        (
            make_index[1.0],
            False,
            ValueError,
        ),
        (
            make_index[1.0, 2.0],
            False,
            ValueError,
        ),
        (
            make_index[1.0, 2.0, 3.0],
            False,
            ValueError,
        ),
        (
            make_index[True],
            False,
            ValueError,
        ),
        (
            make_index[True, False],
            False,
            ValueError,
        ),
        (
            make_index[True, False, True],
            False,
            ValueError,
        ),
        (
            make_index[None],
            False,
            ValueError,
        ),
        (
            make_index[None, None],
            False,
            ValueError,
        ),
        (
            make_index[None, None, None],
            False,
            ValueError,
        ),
        # Range selection
        (
            make_index[:, :],
            False,
            (0, 0, 10, 20),
        ),
        (
            make_index[5:, :],
            False,
            (0, 5, 10, 15),
        ),
        (
            make_index[:5, :],
            False,
            (0, 0, 10, 5),
        ),
        (
            make_index[:, 5:],
            False,
            (5, 0, 5, 20),
        ),
        (
            make_index[:, :5],
            False,
            (0, 0, 5, 20),
        ),
        # Range selection, but negative
        (
            make_index[-2:, :],
            False,
            (0, 18, 10, 2),
        ),
        (
            make_index[:-2, :],
            False,
            (0, 0, 10, 18),
        ),
        (
            make_index[:, -2:],
            False,
            (8, 0, 2, 20),
        ),
        (
            make_index[:, :-2],
            False,
            (0, 0, 8, 20),
        ),
        # Out of bounds
        (
            make_index[-50, :],
            False,
            IndexError,
        ),
        (
            make_index[:, -50],
            False,
            IndexError,
        ),
        # Ranged out of bounds
        (
            make_index[-60:, :],
            False,
            IndexError,
        ),
        (
            make_index[:-60, :],
            False,
            IndexError,
        ),
        (
            make_index[:, -60:],
            False,
            IndexError,
        ),
        (
            make_index[:, :-60],
            False,
            IndexError,
        ),
        # Incorrect types inside of slices
        (
            make_index[2.0:, :],
            False,
            ValueError,
        ),
        (
            make_index[:True, :],
            False,
            ValueError,
        ),
        (
            make_index[:, "Heyho":],
            False,
            ValueError,
        ),
        (
            make_index[:, :list],
            False,
            ValueError,
        ),
        # Indexing into the header
        (
            make_index["header", :],
            False,
            (0, "header", 10, 1),
        ),
        (
            make_index["header", 3:5],
            False,
            (3, "header", 2, 1),
        ),
        (
            make_index["header", -2:],
            False,
            (8, "header", 2, 1),
        ),
        (
            make_index["header", :-2],
            False,
            (0, "header", 8, 1),
        ),
        # Indexing into the header, but in the wrong axis
        (
            make_index[0, "header"],
            False,
            ValueError,
        ),
    ],
)
def test_indices(
    index: t.Any,
    enable_column_names: bool,
    result_should: tuple[
        int,
        int | t.Literal["header"],
        int,
        int,
    ]
    | t.Type[Exception],
) -> None:
    if enable_column_names:
        column_names = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]
    else:
        column_names = None

    # Should this work, or fail?
    if isinstance(result_should, tuple):
        # Let Rio bring the selection into the format (left, top, width, height)
        result_is = rio.components.table._indices_to_rectangle(
            index=index,
            column_names=column_names,
            data_width=10,
            data_height=20,
        )

        # Does reality match expectations?
        assert result_is == result_should

    else:
        with pytest.raises(result_should):
            rio.components.table._indices_to_rectangle(
                index,
                column_names=column_names,
                data_width=10,
                data_height=20,
            )