polars-runtime-compat 1.34.0b2__cp39-abi3-win_amd64.whl → 1.34.0b4__cp39-abi3-win_amd64.whl

polars/testing/parametric/strategies/core.py

@@ -1,615 +0,0 @@
-from __future__ import annotations
-
-from collections.abc import Mapping
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, overload
-
-import hypothesis.strategies as st
-from hypothesis.errors import InvalidArgument
-
-from polars import select, when
-from polars._utils.deprecation import issue_deprecation_warning
-from polars.dataframe import DataFrame
-from polars.datatypes import Array, Boolean, DataType, DataTypeClass, List, Null, Struct
-from polars.series import Series
-from polars.testing.parametric.strategies._utils import flexhash
-from polars.testing.parametric.strategies.data import data
-from polars.testing.parametric.strategies.dtype import _instantiate_dtype, dtypes
-
-if TYPE_CHECKING:
-    from collections.abc import Collection, Sequence
-    from typing import Literal
-
-    from hypothesis.strategies import DrawFn, SearchStrategy
-
-    from polars import LazyFrame
-    from polars._typing import PolarsDataType
-
-
-_ROW_LIMIT = 5  # max generated frame/series length
-_COL_LIMIT = 5  # max number of generated cols
-
-
-@st.composite
-def series(
-    draw: DrawFn,
-    /,
-    *,
-    name: str | SearchStrategy[str] | None = None,
-    dtype: PolarsDataType | None = None,
-    min_size: int = 0,
-    max_size: int = _ROW_LIMIT,
-    strategy: SearchStrategy[Any] | None = None,
-    allow_null: bool = True,
-    allow_chunks: bool = True,
-    allow_masked_out: bool = True,
-    unique: bool = False,
-    allowed_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
-    excluded_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
-    allow_time_zones: bool = True,
-    **kwargs: Any,
-) -> Series:
-    """
-    Hypothesis strategy for producing Polars Series.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    Parameters
-    ----------
-    name : {str, strategy}, optional
-        literal string or a strategy for strings (or None), passed to the Series
-        constructor name-param.
-    dtype : PolarsDataType, optional
-        a valid polars DataType for the resulting series.
-    min_size : int
-        if not passing an exact size, can set a minimum here (defaults to 0).
-        no-op if `size` is set.
-    max_size : int
-        if not passing an exact size, can set a maximum value here (defaults to
-        MAX_DATA_SIZE). no-op if `size` is set.
-    strategy : strategy, optional
-        supports overriding the default strategy for the given dtype.
-    allow_null : bool
-        Allow nulls as possible values and allow the `Null` data type by default.
-    allow_chunks : bool
-        Allow the Series to contain multiple chunks.
-    allow_masked_out : bool
-        Allow the nulls to contain masked out elements.
-    unique : bool, optional
-        indicate whether Series values should all be distinct.
-    allowed_dtypes : {list,set}, optional
-        when automatically generating Series data, allow only these dtypes.
-    excluded_dtypes : {list,set}, optional
-        when automatically generating Series data, exclude these dtypes.
-    allow_time_zones
-        Allow generating `Datetime` Series with a time zone.
-    **kwargs
-        Additional keyword arguments that are passed to the underlying data generation
-        strategies.
-
-    size : int, optional
-        if set, creates a Series of exactly this size (ignoring min_size/max_size
-        params).
-
-        .. deprecated:: 1.0.0
-            Use `min_size` and `max_size` instead.
-
-    null_probability : float
-        Percentage chance (expressed between 0.0 => 1.0) that any Series value is null.
-        This is applied independently of any None values generated by the underlying
-        strategy.
-
-        .. deprecated:: 0.20.26
-            Use `allow_null` instead.
-
-    allow_infinities : bool, optional
-        Allow generation of +/-inf values for floating-point dtypes.
-
-        .. deprecated:: 0.20.26
-            Use `allow_infinity` instead.
-
-    Notes
-    -----
-    In actual usage this is deployed as a unit test decorator, providing a strategy
-    that generates multiple Series with the given dtype/size characteristics for the
-    unit test. While developing a strategy/test, it can also be useful to call
-    `.example()` directly on a given strategy to see concrete instances of the
-    generated data.
-
-    Examples
-    --------
-    The strategy is generally used to generate series in a unit test:
-
-    >>> from polars.testing.parametric import series
-    >>> from hypothesis import given
-    >>> @given(s=series(min_size=3, max_size=5))
-    ... def test_series_len(s: pl.Series) -> None:
-    ...     assert 3 <= s.len() <= 5
-
-    Drawing examples interactively is also possible with the `.example()` method.
-    This should be avoided while running tests.
-
-    >>> from polars.testing.parametric import lists
-    >>> s = series(strategy=lists(pl.String, select_from=["xx", "yy", "zz"]))
-    >>> s.example()  # doctest: +SKIP
-    shape: (4,)
-    Series: '' [list[str]]
-    [
-            ["zz", "zz"]
-            ["zz", "xx", "yy"]
-            []
-            ["xx"]
-    ]
-    """
-    if (null_prob := kwargs.pop("null_probability", None)) is not None:
-        allow_null = _handle_null_probability_deprecation(null_prob)  # type: ignore[assignment]
-    if (allow_inf := kwargs.pop("allow_infinities", None)) is not None:
-        issue_deprecation_warning(
-            "`allow_infinities` is deprecated. Use `allow_infinity` instead.",
-            version="0.20.26",
-        )
-        kwargs["allow_infinity"] = allow_inf
-    if (chunked := kwargs.pop("chunked", None)) is not None:
-        issue_deprecation_warning(
-            "`chunked` is deprecated. Use `allow_chunks` instead.",
-            version="0.20.26",
-        )
-        allow_chunks = chunked
-    if (size := kwargs.pop("size", None)) is not None:
-        issue_deprecation_warning(
-            "`size` is deprecated. Use `min_size` and `max_size` instead.",
-            version="1.0.0",
-        )
-        min_size = max_size = size
-
-    if isinstance(allowed_dtypes, (DataType, DataTypeClass)):
-        allowed_dtypes = [allowed_dtypes]
-    elif allowed_dtypes is not None:
-        allowed_dtypes = list(allowed_dtypes)
-    if isinstance(excluded_dtypes, (DataType, DataTypeClass)):
-        excluded_dtypes = [excluded_dtypes]
-    elif excluded_dtypes is not None:
-        excluded_dtypes = list(excluded_dtypes)
-
-    if not allow_null and not (allowed_dtypes is not None and Null in allowed_dtypes):
-        if excluded_dtypes is None:
-            excluded_dtypes = [Null]
-        else:
-            excluded_dtypes.append(Null)
-
-    if strategy is None:
-        if dtype is None:
-            dtype_strat = dtypes(
-                allowed_dtypes=allowed_dtypes,
-                excluded_dtypes=excluded_dtypes,
-                allow_time_zones=allow_time_zones,
-            )
-        else:
-            dtype_strat = _instantiate_dtype(
-                dtype,
-                allowed_dtypes=allowed_dtypes,
-                excluded_dtypes=excluded_dtypes,
-                allow_time_zones=allow_time_zones,
-            )
-        dtype = draw(dtype_strat)
-
-    if min_size == max_size:
-        size = min_size
-    else:
-        size = draw(st.integers(min_value=min_size, max_value=max_size))
-
-    if isinstance(name, st.SearchStrategy):
-        name = draw(name)
-
-    do_mask_out = (
-        allow_masked_out
-        and allow_null
-        and isinstance(dtype, (List, Array, Struct))
-        and draw(st.booleans())
-    )
-
-    if size == 0:
-        values = []
-    else:
-        # Create series using dtype-specific strategy to generate values
-        if strategy is None:
-            strategy = data(
-                dtype,  # type: ignore[arg-type]
-                allow_null=allow_null and not do_mask_out,
-                **kwargs,
-            )
-
-        values = draw(
-            st.lists(
-                strategy,
-                min_size=size,
-                max_size=size,
-                unique_by=(flexhash if unique else None),
-            )
-        )
-
-    s = Series(name=name, values=values, dtype=dtype)
-
-    # Apply masking out of values
-    if do_mask_out:
-        values = draw(
-            st.lists(
-                st.booleans(),
-                min_size=size,
-                max_size=size,
-                unique_by=(flexhash if unique else None),
-            )
-        )
-
-        mask = Series(name=None, values=values, dtype=Boolean)
-        s = select(when(mask).then(s).alias(s.name)).to_series()
-
-    # Apply chunking
-    if allow_chunks and size > 1 and draw(st.booleans()):
-        split_at = size // 2
-        s = s[:split_at].append(s[split_at:])
-
-    return s
-
-
-@overload
-def dataframes(
-    cols: int | column | Sequence[column] | None = None,
-    *,
-    lazy: Literal[False] = ...,
-    min_cols: int = 0,
-    max_cols: int = _COL_LIMIT,
-    min_size: int = 0,
-    max_size: int = _ROW_LIMIT,
-    include_cols: Sequence[column] | column | None = None,
-    allow_null: bool | Mapping[str, bool] = True,
-    allow_chunks: bool = True,
-    allow_masked_out: bool = True,
-    allowed_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
-    excluded_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
-    allow_time_zones: bool = True,
-    **kwargs: Any,
-) -> SearchStrategy[DataFrame]: ...
-
-
-@overload
-def dataframes(
-    cols: int | column | Sequence[column] | None = None,
-    *,
-    lazy: Literal[True],
-    min_cols: int = 0,
-    max_cols: int = _COL_LIMIT,
-    min_size: int = 0,
-    max_size: int = _ROW_LIMIT,
-    include_cols: Sequence[column] | column | None = None,
-    allow_null: bool | Mapping[str, bool] = True,
-    allow_chunks: bool = True,
-    allow_masked_out: bool = True,
-    allowed_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
-    excluded_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
-    allow_time_zones: bool = True,
-    **kwargs: Any,
-) -> SearchStrategy[LazyFrame]: ...
-
-
-@st.composite
-def dataframes(
-    draw: DrawFn,
-    /,
-    cols: int | column | Sequence[column] | None = None,
-    *,
-    lazy: bool = False,
-    min_cols: int = 1,
-    max_cols: int = _COL_LIMIT,
-    min_size: int = 0,
-    max_size: int = _ROW_LIMIT,
-    include_cols: Sequence[column] | column | None = None,
-    allow_null: bool | Mapping[str, bool] = True,
-    allow_chunks: bool = True,
-    allow_masked_out: bool = True,
-    allowed_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
-    excluded_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
-    allow_time_zones: bool = True,
-    **kwargs: Any,
-) -> DataFrame | LazyFrame:
-    """
-    Hypothesis strategy for producing Polars DataFrames or LazyFrames.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    Parameters
-    ----------
-    cols : {int, columns}, optional
-        integer number of columns to create, or a sequence of `column` objects
-        that describe the desired DataFrame column data.
-    lazy : bool, optional
-        produce a LazyFrame instead of a DataFrame.
-    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).
-    min_size : int, optional
-        if not passing an exact size, set the minimum number of rows in the
-        DataFrame.
-    max_size : int, optional
-        if not passing an exact size, set the maximum number of rows in the
-        DataFrame.
-    include_cols : [column], optional
-        a list of `column` objects to include in the generated DataFrame. note that
-        explicitly provided columns are appended onto the list of existing columns
-        (if any present).
-    allow_null : bool or Mapping[str, bool]
-        Allow nulls as possible values and allow the `Null` data type by default.
-        Accepts either a boolean or a mapping of column names to booleans.
-    allow_chunks : bool
-        Allow the DataFrame to contain multiple chunks.
-    allow_masked_out : bool
-        Allow the nulls to contain masked out elements.
-    allowed_dtypes : {list,set}, optional
-        when automatically generating data, allow only these dtypes.
-    excluded_dtypes : {list,set}, optional
-        when automatically generating data, exclude these dtypes.
-    allow_time_zones
-        Allow generating `Datetime` columns with a time zone.
-    **kwargs
-        Additional keyword arguments that are passed to the underlying data generation
-        strategies.
-
-    size : int, optional
-        if set, will create a DataFrame of exactly this size (and ignore
-        the min_size/max_size len params).
-
-        .. deprecated:: 1.0.0
-            Use `min_size` and `max_size` instead.
-
-    null_probability : {float, dict[str,float]}, optional
-        percentage chance (expressed between 0.0 => 1.0) that a generated value is
-        None. this is applied independently of any None values generated by the
-        underlying strategy, and can be applied either on a per-column basis (if
-        given as a `{col:pct}` dict), or globally. if null_probability is defined
-        on a column, it takes precedence over the global value.
-
-        .. deprecated:: 0.20.26
-            Use `allow_null` instead.
-
-    allow_infinities : bool, optional
-        optionally disallow generation of +/-inf values for floating-point dtypes.
-
-        .. deprecated:: 0.20.26
-            Use `allow_infinity` instead.
-
-    Notes
-    -----
-    In actual usage this is deployed as a unit test decorator, providing a strategy
-    that generates DataFrames or LazyFrames with the given characteristics for
-    the unit test. While developing a strategy/test, it can also be useful to
-    call `.example()` directly on a given strategy to see concrete instances of
-    the generated data.
-
-    Examples
-    --------
-    The strategy is generally used to generate series in a unit test:
-
-    >>> from polars.testing.parametric import dataframes
-    >>> from hypothesis import given
-    >>> @given(df=dataframes(min_size=3, max_size=5))
-    ... def test_df_height(df: pl.DataFrame) -> None:
-    ...     assert 3 <= df.height <= 5
-
-    Drawing examples interactively is also possible with the `.example()` method.
-    This should be avoided while running tests.
-
-    >>> df = dataframes(allowed_dtypes=[pl.Datetime, pl.Float64], max_cols=3)
-    >>> df.example()  # doctest: +SKIP
-    shape: (3, 3)
-    ┌─────────────┬────────────────────────────┬───────────┐
-    │ col0        ┆ col1                       ┆ col2      │
-    │ ---         ┆ ---                        ┆ ---       │
-    │ f64         ┆ datetime[ns]               ┆ f64       │
-    ╞═════════════╪════════════════════════════╪═══════════╡
-    │ NaN         ┆ 1844-07-05 06:19:48.848808 ┆ 3.1436e16 │
-    │ -1.9914e218 ┆ 2068-12-01 23:05:11.412277 ┆ 2.7415e16 │
-    │ 0.5         ┆ 2095-11-19 22:05:17.647961 ┆ -0.5      │
-    └─────────────┴────────────────────────────┴───────────┘
-
-    Use :class:`column` for more control over which exactly which columns are generated.
-
-    >>> from polars.testing.parametric import column
-    >>> dfs = dataframes(
-    ...     [
-    ...         column("x", dtype=pl.Int32),
-    ...         column("y", dtype=pl.Float64),
-    ...     ],
-    ...     min_size=2,
-    ...     max_size=2,
-    ... )
-    >>> dfs.example()  # doctest: +SKIP
-    shape: (2, 2)
-    ┌───────────┬────────────┐
-    │ x         ┆ y          │
-    │ ---       ┆ ---        │
-    │ i32       ┆ f64        │
-    ╞═══════════╪════════════╡
-    │ -15836    ┆ 1.1755e-38 │
-    │ 575050513 ┆ NaN        │
-    └───────────┴────────────┘
-    """
-    if (null_prob := kwargs.pop("null_probability", None)) is not None:
-        allow_null = _handle_null_probability_deprecation(null_prob)
-    if (allow_inf := kwargs.pop("allow_infinities", None)) is not None:
-        issue_deprecation_warning(
-            "`allow_infinities` is deprecated. Use `allow_infinity` instead.",
-            version="0.20.26",
-        )
-        kwargs["allow_infinity"] = allow_inf
-    if (chunked := kwargs.pop("chunked", None)) is not None:
-        issue_deprecation_warning(
-            "`chunked` is deprecated. Use `allow_chunks` instead.",
-            version="0.20.26",
-        )
-        allow_chunks = chunked
-    if (size := kwargs.pop("size", None)) is not None:
-        issue_deprecation_warning(
-            "`size` is deprecated. Use `min_size` and `max_size` instead.",
-            version="1.0.0",
-        )
-        min_size = max_size = size
-
-    if isinstance(include_cols, column):
-        include_cols = [include_cols]
-
-    if cols is None:
-        n_cols = draw(st.integers(min_value=min_cols, max_value=max_cols))
-        cols = [column() for _ in range(n_cols)]
-    elif isinstance(cols, int):
-        cols = [column() for _ in range(cols)]
-    elif isinstance(cols, column):
-        cols = [cols]
-    else:
-        cols = list(cols)
-
-    if include_cols:
-        cols.extend(list(include_cols))
-
-    if min_size == max_size:
-        size = min_size
-    else:
-        size = draw(st.integers(min_value=min_size, max_value=max_size))
-
-    # Process columns
-    for idx, c in enumerate(cols):
-        if c.name is None:
-            c.name = f"col{idx}"
-        if c.allow_null is None:
-            if isinstance(allow_null, Mapping):
-                c.allow_null = allow_null.get(c.name, True)
-            else:
-                c.allow_null = allow_null
-
-    allow_series_chunks = draw(st.booleans()) if allow_chunks else False
-
-    data = {
-        c.name: draw(
-            series(
-                name=c.name,
-                dtype=c.dtype,
-                min_size=size,
-                max_size=size,
-                strategy=c.strategy,
-                allow_null=c.allow_null,  # type: ignore[arg-type]
-                allow_chunks=allow_series_chunks,
-                allow_masked_out=allow_masked_out,
-                unique=c.unique,
-                allowed_dtypes=allowed_dtypes,
-                excluded_dtypes=excluded_dtypes,
-                allow_time_zones=allow_time_zones,
-                **kwargs,
-            )
-        )
-        for c in cols
-    }
-
-    df = DataFrame(data)
-
-    # Apply chunking
-    if allow_chunks and size > 1 and not allow_series_chunks and draw(st.booleans()):
-        split_at = size // 2
-        df = df[:split_at].vstack(df[split_at:])
-
-    if lazy:
-        return df.lazy()
-
-    return df
-
-
-@dataclass
-class column:
-    """
-    Define a column for use with the `dataframes` strategy.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    Parameters
-    ----------
-    name : str
-        string column name.
-    dtype : PolarsDataType
-        a polars dtype.
-    strategy : strategy, optional
-        supports overriding the default strategy for the given dtype.
-    allow_null : bool, optional
-        Allow nulls as possible values and allow the `Null` data type by default.
-    unique : bool, optional
-        flag indicating that all values generated for the column should be unique.
-
-    null_probability : float, optional
-        percentage chance (expressed between 0.0 => 1.0) that a generated value is
-        None. this is applied independently of any None values generated by the
-        underlying strategy.
-
-        .. deprecated:: 0.20.26
-            Use `allow_null` instead.
-
-    Examples
-    --------
-    >>> from polars.testing.parametric import column
-    >>> dfs = dataframes(
-    ...     [
-    ...         column("x", dtype=pl.Int32, allow_null=True),
-    ...         column("y", dtype=pl.Float64),
-    ...     ],
-    ...     size=2,
-    ... )
-    >>> dfs.example()  # doctest: +SKIP
-    shape: (2, 2)
-    ┌───────────┬────────────┐
-    │ x         ┆ y          │
-    │ ---       ┆ ---        │
-    │ i32       ┆ f64        │
-    ╞═══════════╪════════════╡
-    │ null      ┆ 1.1755e-38 │
-    │ 575050513 ┆ inf        │
-    └───────────┴────────────┘
-    """
-
-    name: str | None = None
-    dtype: PolarsDataType | None = None
-    strategy: SearchStrategy[Any] | None = None
-    allow_null: bool | None = None
-    unique: bool = False
-
-    null_probability: float | None = None
-
-    def __post_init__(self) -> None:
-        if self.null_probability is not None:
-            self.allow_null = _handle_null_probability_deprecation(  # type: ignore[assignment]
-                self.null_probability
-            )
-
-
-def _handle_null_probability_deprecation(
-    null_probability: float | Mapping[str, float],
-) -> bool | dict[str, bool]:
-    issue_deprecation_warning(
-        "`null_probability` is deprecated. Use `allow_null` instead.",
-        version="0.20.26",
-    )
-
-    def prob_to_bool(prob: float) -> bool:
-        if not (0.0 <= prob <= 1.0):
-            msg = f"`null_probability` should be between 0.0 and 1.0, got {prob!r}"
-            raise InvalidArgument(msg)
-
-        return bool(prob)
-
-    if isinstance(null_probability, Mapping):
-        return {col: prob_to_bool(prob) for col, prob in null_probability.items()}
-    else:
-        return prob_to_bool(null_probability)