Repository: pola-rs/polars
File: py-polars/tests/unit/operations/namespaces/test_categorical.py (Python)

Code before:

def test_sort_categoricals_6014_lexical() -> None:
    # create lexically-ordered categorical
    df = pl.DataFrame({"key": ["bbb", "aaa", "ccc"]}).with_columns(
        pl.col("key").cast(pl.Categorical())
    )

    out = df.sort("key")
    assert out.to_dict(as_series=False) == {"key": ["aaa", "bbb", "ccc"]}


def test_categorical_get_categories() -> None:
    s = pl.Series("cats", ["foo", "bar", "foo", "foo", "ham"], dtype=pl.Categorical)
    assert set(s.cat.get_categories().to_list()) >= {"foo", "bar", "ham"}


def test_cat_to_local() -> None:
    s = pl.Series(["a", "b", "a"], dtype=pl.Categorical)
    assert_series_equal(s, s.cat.to_local())


def test_cat_uses_lexical_ordering() -> None:
    with pytest.warns(DeprecationWarning, match="ordering parameter"):
        physical_cat = pl.Categorical(ordering="physical")

    for dtype in [pl.Categorical, pl.Categorical(), physical_cat]:
        s = pl.Series(["a", "b", None, "b"]).cast(dtype)

        with pytest.warns(
            DeprecationWarning,
            match="Categoricals are now always ordered lexically",
        ):
            assert s.cat.uses_lexical_ordering()


@pytest.mark.parametrize("dtype", [pl.Categorical, pl.Enum])
def test_cat_len_bytes(dtype: PolarsDataType) -> None:
    # test Series
    values = ["Café", None, "Café", "345", "東京"]
    if dtype == pl.Enum:
        dtype = pl.Enum(list({x for x in values if x is not None}))
    s = pl.Series("a", values, dtype=dtype)
    result = s.cat.len_bytes()
    expected = pl.Series("a", [5, None, 5, 3, 6], dtype=pl.UInt32)
    assert_series_equal(result, expected)

    # test DataFrame expr
    df = pl.DataFrame(s)
    result_df = df.select(pl.col("a").cat.len_bytes())
    expected_df = pl.DataFrame(expected)
    assert_frame_equal(result_df, expected_df)

Review comment (suggestion):
Iterate over different variants of the Categorical data type to reduce duplication (since Ruff requires pytest.warns to only wrap one statement I iterate over the types, instead of adding 3x pytest.warns)

Code after:

def test_sort_categoricals_6014_lexical() -> None:
    # create lexically-ordered categorical
    df = pl.DataFrame({"key": ["bbb", "aaa", "ccc"]}).with_columns(
        pl.col("key").cast(pl.Categorical())
    )

    out = df.sort("key")
    assert out.to_dict(as_series=False) == {"key": ["aaa", "bbb", "ccc"]}


def test_categorical_get_categories() -> None:
    s = pl.Series("cats", ["foo", "bar", "foo", "foo", "ham"], dtype=pl.Categorical)
    assert set(s.cat.get_categories().to_list()) >= {"foo", "bar", "ham"}


def test_cat_to_local() -> None:
    s = pl.Series(["a", "b", "a"], dtype=pl.Categorical)
    assert_series_equal(s, s.cat.to_local())


def test_cat_uses_lexical_ordering() -> None:
    with pytest.warns(DeprecationWarning, match="ordering parameter"):
        physical_cat = pl.Categorical(ordering="physical")

    for dtype in [pl.Categorical, pl.Categorical(), physical_cat]:
        s = pl.Series(["a", "b", None, "b"]).cast(dtype)  # type: ignore[arg-type]

        with pytest.warns(
            DeprecationWarning,
            match="Categoricals are now always ordered lexically",
        ):
            assert s.cat.uses_lexical_ordering()


@pytest.mark.parametrize("dtype", [pl.Categorical, pl.Enum])
def test_cat_len_bytes(dtype: PolarsDataType) -> None:
    # test Series
    values = ["Café", None, "Café", "345", "東京"]
    if dtype == pl.Enum:
        dtype = pl.Enum(list({x for x in values if x is not None}))
    s = pl.Series("a", values, dtype=dtype)
    result = s.cat.len_bytes()
    expected = pl.Series("a", [5, None, 5, 3, 6], dtype=pl.UInt32)
    assert_series_equal(result, expected)

    # test DataFrame expr
    df = pl.DataFrame(s)
    result_df = df.select(pl.col("a").cat.len_bytes())
    expected_df = pl.DataFrame(expected)
    assert_frame_equal(result_df, expected_df)