Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
 Path: blob/main/py-polars/tests/unit/testing/parametric/strategies/test_data.py
8424 views
1
from __future__ import annotations

2


3
from hypothesis import given

4


5
import polars as pl

6
from polars.testing.parametric.strategies.data import categories, data

7


8


9
@given(cat=categories(3))

10
def test_categories(cat: str) -> None:

11
    assert cat in ("c0", "c1", "c2")

12


13


14
@given(cat=data(pl.Categorical, n_categories=3))

15
def test_data_kwargs(cat: str) -> None:

16
    assert cat in ("c0", "c1", "c2")

17


18


19
@given(categories=data(pl.List(pl.Categorical), n_categories=3))

20
def test_data_nested_kwargs(categories: list[str]) -> None:

21
    assert all(c in ("c0", "c1", "c2") for c in categories)

22


23


24
@given(cat=data(pl.Enum))

25
def test_data_enum(cat: str) -> None:

26
    assert cat in ("c0", "c1", "c2")

27


28


29
@given(cat=data(pl.Enum(["hello", "world"])))

30
def test_data_enum_instantiated(cat: str) -> None:

31
    assert cat in ("hello", "world")

32


33


34
@given(struct=data(pl.Struct({"a": pl.Int8, "b": pl.String})))

35
def test_data_struct(struct: dict[str, int | str]) -> None:

36
    assert isinstance(struct["a"], int)

37
    assert isinstance(struct["b"], str)

38


39