polars-runtime-compat 1.34.0b2__cp39-abi3-macosx_10_12_x86_64.whl

polars/testing/parametric/strategies/dtype.py

@@ -0,0 +1,436 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import hypothesis.strategies as st
+from hypothesis.errors import InvalidArgument
+
+from polars.datatypes import (
+    Array,
+    Binary,
+    Boolean,
+    Categorical,
+    DataType,
+    Date,
+    Datetime,
+    Decimal,
+    Duration,
+    Enum,
+    Field,
+    Float32,
+    Float64,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Int128,
+    List,
+    Null,
+    String,
+    Struct,
+    Time,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+    UInt128,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Collection, Sequence
+
+    from hypothesis.strategies import DrawFn, SearchStrategy
+
+    from polars._typing import PolarsDataType, TimeUnit
+    from polars.datatypes import DataTypeClass
+
+
+# Supported data type classes which do not take any arguments
+_SIMPLE_DTYPES: list[DataTypeClass] = [
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Int128,
+    Float64,
+    Float32,
+    Boolean,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+    UInt128,
+    String,
+    Binary,
+    Date,
+    Time,
+    Null,
+    # TODO: Enable Object types by default when various issues are solved.
+    # Object,
+]
+# Supported data type classes with arguments
+_COMPLEX_DTYPES: list[DataTypeClass] = [
+    Datetime,
+    Duration,
+    Categorical,
+    Decimal,
+    Enum,
+]
+# Supported data type classes that contain other data types
+_NESTED_DTYPES: list[DataTypeClass] = [
+    # TODO: Enable nested types by default when various issues are solved.
+    # List,
+    # Array,
+    Struct,
+]
+# Supported data type classes that do not contain other data types
+_FLAT_DTYPES = _SIMPLE_DTYPES + _COMPLEX_DTYPES
+
+_DEFAULT_ARRAY_WIDTH_LIMIT = 3
+_DEFAULT_STRUCT_FIELDS_LIMIT = 3
+_DEFAULT_ENUM_CATEGORIES_LIMIT = 3
+
+
+def dtypes(
+    *,
+    allowed_dtypes: Collection[PolarsDataType] | None = None,
+    excluded_dtypes: Sequence[PolarsDataType] | None = None,
+    allow_time_zones: bool = True,
+    nesting_level: int = 3,
+) -> SearchStrategy[DataType]:
+    """
+    Create a strategy for generating Polars :class:`DataType` objects.
+
+    .. warning::
+        This functionality is currently considered **unstable**. It may be
+        changed at any point without it being considered a breaking change.
+
+    Parameters
+    ----------
+    allowed_dtypes
+        Data types the strategy will pick from. If set to `None` (default),
+        all supported data types are included.
+    excluded_dtypes
+        Data types the strategy will *not* pick from. This takes priority over
+        data types specified in `allowed_dtypes`.
+    allow_time_zones
+        Allow generating `Datetime` data types with a time zone.
+    nesting_level
+        The complexity of nested data types. If set to 0, nested data types are
+        disabled.
+    """
+    flat_dtypes, nested_dtypes, excluded_dtypes = _parse_dtype_restrictions(
+        allowed_dtypes, excluded_dtypes
+    )
+
+    if nesting_level > 0 and nested_dtypes:
+        if not flat_dtypes:
+            return _nested_dtypes(
+                inner=st.just(Null()),
+                allowed_dtypes=nested_dtypes,
+                excluded_dtypes=excluded_dtypes,
+                allow_time_zones=allow_time_zones,
+            )
+        return st.recursive(
+            base=_flat_dtypes(
+                allowed_dtypes=flat_dtypes,
+                excluded_dtypes=excluded_dtypes,
+                allow_time_zones=allow_time_zones,
+            ),
+            extend=lambda s: _nested_dtypes(
+                s,
+                allowed_dtypes=nested_dtypes,
+                excluded_dtypes=excluded_dtypes,
+                allow_time_zones=allow_time_zones,
+            ),
+            max_leaves=nesting_level,
+        )
+    else:
+        return _flat_dtypes(
+            allowed_dtypes=flat_dtypes,
+            excluded_dtypes=excluded_dtypes,
+            allow_time_zones=allow_time_zones,
+        )
+
+
+def _parse_dtype_restrictions(
+    allowed_dtypes: Collection[PolarsDataType] | None = None,
+    excluded_dtypes: Sequence[PolarsDataType] | None = None,
+) -> tuple[list[PolarsDataType], list[PolarsDataType], list[DataType]]:
+    """
+    Parse data type restrictions.
+
+    Splits allowed data types into flat and nested data types.
+    Filters the allowed data types by excluded data type classes.
+    Excluded instantiated data types are returned to be filtered later.
+    """
+    # Split excluded dtypes into instances and classes
+    excluded_dtypes_instance = []
+    excluded_dtypes_class = []
+    if excluded_dtypes:
+        for dt in excluded_dtypes:
+            if isinstance(dt, DataType):
+                excluded_dtypes_instance.append(dt)
+            else:
+                excluded_dtypes_class.append(dt)
+
+    # Split allowed dtypes into flat and nested, excluding certain dtype classes
+    allowed_dtypes_flat: list[PolarsDataType]
+    allowed_dtypes_nested: list[PolarsDataType]
+    if allowed_dtypes is None:
+        allowed_dtypes_flat = [
+            dt for dt in _FLAT_DTYPES if dt not in excluded_dtypes_class
+        ]
+        allowed_dtypes_nested = [
+            dt for dt in _NESTED_DTYPES if dt not in excluded_dtypes_class
+        ]
+    else:
+        allowed_dtypes_flat = []
+        allowed_dtypes_nested = []
+        for dt in allowed_dtypes:
+            if dt in excluded_dtypes_class:
+                continue
+            elif dt.is_nested():
+                allowed_dtypes_nested.append(dt)
+            else:
+                allowed_dtypes_flat.append(dt)
+
+    return allowed_dtypes_flat, allowed_dtypes_nested, excluded_dtypes_instance
+
+
+@st.composite
+def _flat_dtypes(
+    draw: DrawFn,
+    allowed_dtypes: Sequence[PolarsDataType] | None = None,
+    excluded_dtypes: Sequence[PolarsDataType] | None = None,
+    *,
+    allow_time_zones: bool = True,
+) -> DataType:
+    """Create a strategy for generating non-nested Polars :class:`DataType` objects."""
+    if allowed_dtypes is None:
+        allowed_dtypes = _FLAT_DTYPES
+    if excluded_dtypes is None:
+        excluded_dtypes = []
+
+    dtype = draw(st.sampled_from(allowed_dtypes))
+    return draw(
+        _instantiate_flat_dtype(dtype, allow_time_zones=allow_time_zones).filter(
+            lambda x: x not in excluded_dtypes
+        )
+    )
+
+
+@st.composite
+def _instantiate_flat_dtype(
+    draw: DrawFn, dtype: PolarsDataType, *, allow_time_zones: bool = True
+) -> DataType:
+    """Take a flat data type and instantiate it."""
+    if isinstance(dtype, DataType):
+        return dtype
+    elif dtype in _SIMPLE_DTYPES:
+        return dtype()
+    elif dtype == Datetime:
+        time_unit = draw(_time_units())
+        time_zone = draw(st.none() | _time_zones()) if allow_time_zones else None
+        return Datetime(time_unit, time_zone)
+    elif dtype == Duration:
+        time_unit = draw(_time_units())
+        return Duration(time_unit)
+    elif dtype == Categorical:
+        return Categorical()
+    elif dtype == Enum:
+        n_categories = draw(
+            st.integers(min_value=1, max_value=_DEFAULT_ENUM_CATEGORIES_LIMIT)
+        )
+        categories = [f"c{i}" for i in range(n_categories)]
+        return Enum(categories)
+    elif dtype == Decimal:
+        precision = draw(st.integers(min_value=1, max_value=38) | st.none())
+        scale = draw(st.integers(min_value=0, max_value=precision or 38))
+        return Decimal(precision, scale)
+    else:
+        msg = f"unsupported data type: {dtype}"
+        raise InvalidArgument(msg)
+
+
+@st.composite
+def _nested_dtypes(
+    draw: DrawFn,
+    inner: SearchStrategy[DataType],
+    allowed_dtypes: Sequence[PolarsDataType] | None = None,
+    excluded_dtypes: Sequence[PolarsDataType] | None = None,
+    *,
+    allow_time_zones: bool = True,
+) -> DataType:
+    """Create a strategy for generating nested Polars :class:`DataType` objects."""
+    if allowed_dtypes is None:
+        allowed_dtypes = _NESTED_DTYPES
+    if excluded_dtypes is None:
+        excluded_dtypes = []
+
+    dtype = draw(st.sampled_from(allowed_dtypes))
+    return draw(
+        _instantiate_nested_dtype(
+            dtype, inner, allow_time_zones=allow_time_zones
+        ).filter(lambda x: x not in excluded_dtypes)
+    )
+
+
+@st.composite
+def _instantiate_nested_dtype(
+    draw: DrawFn,
+    dtype: PolarsDataType,
+    inner: SearchStrategy[DataType],
+    *,
+    allow_time_zones: bool = True,
+) -> DataType:
+    """Take a nested data type and instantiate it."""
+
+    def instantiate_inner(inner_dtype: PolarsDataType | None) -> DataType:
+        if inner_dtype is None:
+            return draw(inner)
+        elif inner_dtype.is_nested():
+            return draw(
+                _instantiate_nested_dtype(
+                    inner_dtype, inner, allow_time_zones=allow_time_zones
+                )
+            )
+        else:
+            return draw(
+                _instantiate_flat_dtype(inner_dtype, allow_time_zones=allow_time_zones)
+            )
+
+    if dtype == List:
+        inner_dtype = instantiate_inner(getattr(dtype, "inner", None))
+        return List(inner_dtype)
+    elif dtype == Array:
+        inner_dtype = instantiate_inner(getattr(dtype, "inner", None))
+        size = getattr(
+            dtype,
+            "size",
+            draw(st.integers(min_value=1, max_value=_DEFAULT_ARRAY_WIDTH_LIMIT)),
+        )
+        return Array(inner_dtype, size)
+    elif dtype == Struct:
+        if isinstance(dtype, Struct):
+            fields = [Field(f.name, instantiate_inner(f.dtype)) for f in dtype.fields]
+        else:
+            n_fields = draw(
+                st.integers(min_value=1, max_value=_DEFAULT_STRUCT_FIELDS_LIMIT)
+            )
+            fields = [Field(f"f{i}", draw(inner)) for i in range(n_fields)]
+        return Struct(fields)
+    else:
+        msg = f"unsupported data type: {dtype}"
+        raise InvalidArgument(msg)
+
+
+def _time_units() -> SearchStrategy[TimeUnit]:
+    """Create a strategy for generating valid units of time."""
+    return st.sampled_from(["us", "ns", "ms"])
+
+
+def _time_zones() -> SearchStrategy[str]:
+    """Create a strategy for generating valid time zones."""
+    # Not available when building docs, so just import here.
+    from polars._plr import _known_timezones
+
+    chrono_known_tz = set(_known_timezones())
+    return st.timezone_keys(allow_prefix=False).filter(
+        lambda tz: tz not in {"Factory", "localtime"} and tz in chrono_known_tz
+    )
+
+
+@st.composite
+def _instantiate_dtype(
+    draw: DrawFn,
+    dtype: PolarsDataType,
+    *,
+    allowed_dtypes: Collection[PolarsDataType] | None = None,
+    excluded_dtypes: Sequence[PolarsDataType] | None = None,
+    nesting_level: int = 3,
+    allow_time_zones: bool = True,
+) -> DataType:
+    """Take a data type and instantiate it."""
+    if not dtype.is_nested():
+        if isinstance(dtype, DataType):
+            return dtype
+
+        if allowed_dtypes is None:
+            allowed_dtypes = [dtype]
+        else:
+            same_dtypes = [dt for dt in allowed_dtypes if dt == dtype]
+            allowed_dtypes = same_dtypes if same_dtypes else [dtype]
+
+        return draw(
+            _flat_dtypes(
+                allowed_dtypes=allowed_dtypes,
+                excluded_dtypes=excluded_dtypes,
+                allow_time_zones=allow_time_zones,
+            )
+        )
+
+    def draw_inner(dtype: PolarsDataType | None) -> DataType:
+        if dtype is None:
+            return draw(
+                dtypes(
+                    allowed_dtypes=allowed_dtypes,
+                    excluded_dtypes=excluded_dtypes,
+                    nesting_level=nesting_level - 1,
+                    allow_time_zones=allow_time_zones,
+                )
+            )
+        else:
+            return draw(
+                _instantiate_dtype(
+                    dtype,
+                    allowed_dtypes=allowed_dtypes,
+                    excluded_dtypes=excluded_dtypes,
+                    nesting_level=nesting_level - 1,
+                    allow_time_zones=allow_time_zones,
+                )
+            )
+
+    if dtype == List:
+        inner = draw_inner(getattr(dtype, "inner", None))
+        return List(inner)
+    elif dtype == Array:
+        inner = draw_inner(getattr(dtype, "inner", None))
+        size = getattr(
+            dtype,
+            "size",
+            draw(st.integers(min_value=1, max_value=_DEFAULT_ARRAY_WIDTH_LIMIT)),
+        )
+        return Array(inner, size)
+    elif dtype == Struct:
+        if isinstance(dtype, Struct):
+            fields = [
+                Field(
+                    name=f.name,
+                    dtype=draw(
+                        _instantiate_dtype(
+                            f.dtype,
+                            allowed_dtypes=allowed_dtypes,
+                            excluded_dtypes=excluded_dtypes,
+                            nesting_level=nesting_level - 1,
+                            allow_time_zones=allow_time_zones,
+                        )
+                    ),
+                )
+                for f in dtype.fields
+            ]
+        else:
+            n_fields = draw(
+                st.integers(min_value=1, max_value=_DEFAULT_STRUCT_FIELDS_LIMIT)
+            )
+            inner_strategy = dtypes(
+                allowed_dtypes=allowed_dtypes,
+                excluded_dtypes=excluded_dtypes,
+                nesting_level=nesting_level - 1,
+                allow_time_zones=allow_time_zones,
+            )
+            fields = [Field(f"f{i}", draw(inner_strategy)) for i in range(n_fields)]
+        return Struct(fields)
+    else:
+        msg = f"unsupported data type: {dtype}"
+        raise InvalidArgument(msg)

polars/testing/parametric/strategies/legacy.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+import hypothesis.strategies as st
+from hypothesis.errors import InvalidArgument
+
+from polars._utils.deprecation import deprecated
+from polars.datatypes import is_polars_dtype
+from polars.testing.parametric.strategies.core import _COL_LIMIT, column
+from polars.testing.parametric.strategies.data import lists
+from polars.testing.parametric.strategies.dtype import _instantiate_dtype, dtypes
+
+if TYPE_CHECKING:
+    import sys
+
+    from hypothesis.strategies import SearchStrategy
+
+    from polars._typing import OneOrMoreDataTypes, PolarsDataType
+
+    if sys.version_info >= (3, 13):
+        from warnings import deprecated
+    else:
+        from typing_extensions import deprecated  # noqa: TC004
+
+
+@deprecated(
+    "`columns` is deprecated; use `column` instead, "
+    "in conjunction with a list comprehension."
+)
+def columns(
+    cols: int | Sequence[str] | None = None,
+    *,
+    dtype: OneOrMoreDataTypes | None = None,
+    min_cols: int = 0,
+    max_cols: int = _COL_LIMIT,
+    unique: bool = False,
+) -> list[column]:
+    """
+    Define multiple columns for use with the @dataframes strategy.
+
+    .. deprecated:: 0.20.26
+        Use :class:`column` instead, in conjunction with a list comprehension.
+
+    .. warning::
+        This functionality is currently considered **unstable**. It may be
+        changed at any point without it being considered a breaking change.
+
+    Generate a fixed sequence of `column` objects suitable for passing to the
+    @dataframes strategy, or using standalone (note that this function is not itself
+    a strategy).
+
+    Notes
+    -----
+    Additional control is available by creating a sequence of columns explicitly,
+    using the `column` class (an especially useful option is to override the default
+    data-generating strategy for a given col/dtype).
+
+    Parameters
+    ----------
+    cols : {int, [str]}, optional
+        integer number of cols to create, or explicit list of column names. if
+        omitted a random number of columns (between mincol and max_cols) are
+        created.
+    dtype : PolarsDataType, optional
+        a single dtype for all cols, or list of dtypes (the same length as `cols`).
+        if omitted, each generated column is assigned a random dtype.
+    min_cols : int, optional
+        if not passing an exact size, can set a minimum here (defaults to 0).
+    max_cols : int, optional
+        if not passing an exact size, can set a maximum value here (defaults to
+        MAX_COLS).
+    unique : bool, optional
+        indicate if the values generated for these columns should be unique
+        (per-column).
+
+    Examples
+    --------
+    >>> from polars.testing.parametric import columns, dataframes
+    >>> from hypothesis import given
+    >>> @given(dataframes(columns(["x", "y", "z"], unique=True)))  # doctest: +SKIP
+    ... def test_unique_xyz(df: pl.DataFrame) -> None:
+    ...     assert_something(df)
+    """
+    # create/assign named columns
+    if cols is None:
+        cols = st.integers(min_value=min_cols, max_value=max_cols).example()
+    if isinstance(cols, int):
+        names: Sequence[str] = [f"col{n}" for n in range(cols)]
+    else:
+        names = cols
+    n_cols = len(names)
+
+    if dtype is None:
+        dtypes: Sequence[PolarsDataType | None] = [None] * n_cols
+    elif is_polars_dtype(dtype):
+        dtypes = [dtype] * n_cols
+    elif isinstance(dtype, Sequence):
+        if (n_dtypes := len(dtype)) != n_cols:
+            msg = f"given {n_dtypes} dtypes for {n_cols} names"
+            raise InvalidArgument(msg)
+        dtypes = dtype
+    else:
+        msg = f"{dtype!r} is not a valid polars datatype"
+        raise InvalidArgument(msg)
+
+    # init list of named/typed columns
+    return [column(name=nm, dtype=tp, unique=unique) for nm, tp in zip(names, dtypes)]
+
+
+@deprecated("`create_list_strategy` is deprecated; use `lists` instead.")
+def create_list_strategy(
+    inner_dtype: PolarsDataType | None = None,
+    *,
+    select_from: Sequence[Any] | None = None,
+    size: int | None = None,
+    min_size: int = 0,
+    max_size: int | None = None,
+    unique: bool = False,
+) -> SearchStrategy[list[Any]]:
+    """
+    Create a strategy for generating Polars :class:`List` data.
+
+    .. deprecated:: 0.20.26
+        Use :func:`lists` instead.
+
+    Parameters
+    ----------
+    inner_dtype : PolarsDataType
+        type of the inner list elements (can also be another List).
+    select_from : list, optional
+        randomly select the innermost values from this list (otherwise
+        the default strategy associated with the innermost dtype is used).
+    size : int, optional
+        if set, generated lists will be of exactly this size (and
+        ignore the min_size/max_size params).
+    min_size : int, optional
+        set the minimum size of the generated lists (default: 0 if unset).
+    max_size : int, optional
+        set the maximum size of the generated lists (default: 3 if
+        min_size is unset or zero, otherwise 2x min_size).
+    unique : bool, optional
+        ensure that the generated lists contain unique values.
+
+    Examples
+    --------
+    Create a strategy that generates a list of i32 values:
+
+    >>> from polars.testing.parametric import create_list_strategy
+    >>> lst = create_list_strategy(inner_dtype=pl.Int32)  # doctest: +SKIP
+    >>> lst.example()  # doctest: +SKIP
+    [-11330, 24030, 116]
+    """
+    if size is not None:
+        min_size = max_size = size
+
+    if inner_dtype is None:
+        inner_dtype = dtypes().example()
+    else:
+        inner_dtype = _instantiate_dtype(inner_dtype).example()
+
+    return lists(
+        inner_dtype,
+        select_from=select_from,
+        min_size=min_size,
+        max_size=max_size,
+        unique=unique,
+    )

polars/type_aliases.py

@@ -0,0 +1,24 @@
+"""
+Deprecated module - do not use.
+
+Used to contain private type aliases. These are now in the `polars._typing` module.
+"""
+
+from typing import Any
+
+import polars._typing as plt
+from polars._utils.deprecation import issue_deprecation_warning
+
+
+def __getattr__(name: str) -> Any:
+    if name in dir(plt):
+        issue_deprecation_warning(
+            "the `polars.type_aliases` module was deprecated in version 1.0.0."
+            " The type aliases have moved to the `polars._typing` module to explicitly mark them as private."
+            " Please define your own type aliases, or temporarily import from the `polars._typing` module."
+            " A public `polars.typing` module will be added in the future.",
+        )
+        return getattr(plt, name)
+
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)