dycw-utilities 0.112.10__tar.gz → 0.112.11__tar.gz

{dycw_utilities-0.112.10 → dycw_utilities-0.112.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.112.10
+Version: 0.112.11
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.112.10 → dycw_utilities-0.112.11}/pyproject.toml

@@ -93,7 +93,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.112.10"
+version = "0.112.11"
 
 [project.optional-dependencies]
 test = [
@@ -336,7 +336,7 @@ zzz-test-zoneinfo = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.112.10"
+current_version = "0.112.11"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.112.10 → dycw_utilities-0.112.11}/src/tests/test_polars.py

@@ -5,7 +5,7 @@ import enum
 import itertools
 from dataclasses import dataclass, field
 from enum import auto
-from itertools import chain
+from itertools import chain, repeat
 from math import isfinite, nan
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
@@ -49,6 +49,7 @@ from polars import (
     datetime_range,
     int_range,
     lit,
+    struct,
 )
 from polars._typing import (
     IntoExprColumn,  # pyright: ignore[reportPrivateImportUsage]
@@ -112,6 +113,10 @@ from utilities.polars import (
     _GetSeriesNumberOfDecimalsNotFloatError,
     _InsertBetweenMissingColumnsError,
     _InsertBetweenNonConsecutiveError,
+    _IsNearEventAfterError,
+    _IsNearEventBeforeError,
+    _ReifyExprsEmptyError,
+    _ReifyExprsSeriesNonUniqueError,
     _yield_struct_series_element_remove_nulls,
     ac_halflife,
     acf,
@@ -144,6 +149,7 @@ from utilities.polars import (
     insert_after,
     insert_before,
     insert_between,
+    is_near_event,
     is_not_null_struct_series,
     is_null_struct_series,
     join,
@@ -151,11 +157,13 @@ from utilities.polars import (
     nan_sum_agg,
     nan_sum_cols,
     normal,
+    reify_exprs,
     replace_time_zone,
     set_first_row_as_columns,
     struct_dtype,
     struct_from_dataclass,
     touch,
+    try_reify_expr,
     uniform,
     unique_element,
     week_num,
@@ -1639,6 +1647,102 @@ class TestIntegers:
         assert series.is_between(0, high, closed="left").all()
 
 
+class TestIsNearEvent:
+    df: ClassVar[DataFrame] = DataFrame(
+        data=[
+            (False, False),
+            (False, False),
+            (True, False),
+            (True, False),
+            (False, False),
+            (False, False),
+            (False, False),
+            (False, False),
+            (False, False),
+            (False, True),
+        ],
+        schema={"x": Boolean, "y": Boolean},
+        orient="row",
+    )
+
+    def test_no_exprs(self) -> None:
+        result = self.df.with_columns(is_near_event().alias("z"))["z"]
+        expected = Series(
+            name="z", values=list(repeat(object=False, times=10)), dtype=Boolean
+        )
+        assert_series_equal(result, expected)
+
+    def test_x(self) -> None:
+        result = self.df.with_columns(is_near_event("x").alias("z"))["z"]
+        expected = Series(
+            name="z",
+            values=[False, False, True, True, False, False, False, False, False, False],
+            dtype=Boolean,
+        )
+        assert_series_equal(result, expected)
+
+    def test_y(self) -> None:
+        result = self.df.with_columns(is_near_event("y").alias("z"))["z"]
+        expected = Series(
+            name="z",
+            values=[
+                False,
+                False,
+                False,
+                False,
+                False,
+                False,
+                False,
+                False,
+                False,
+                True,
+            ],
+            dtype=Boolean,
+        )
+        assert_series_equal(result, expected)
+
+    def test_x_before(self) -> None:
+        result = self.df.with_columns(is_near_event("x", before=1).alias("z"))["z"]
+        expected = Series(
+            name="z",
+            values=[False, True, True, True, False, False, False, False, False, False],
+            dtype=Boolean,
+        )
+        assert_series_equal(result, expected)
+
+    def test_x_after(self) -> None:
+        result = self.df.with_columns(is_near_event("x", after=1).alias("z"))["z"]
+        expected = Series(
+            name="z",
+            values=[False, False, True, True, True, False, False, False, False, False],
+            dtype=Boolean,
+        )
+        assert_series_equal(result, expected)
+
+    def test_x_or_y(self) -> None:
+        result = self.df.with_columns(is_near_event("x", "y").alias("z"))["z"]
+        expected = Series(
+            name="z",
+            values=[False, False, True, True, False, False, False, False, False, True],
+            dtype=Boolean,
+        )
+        assert_series_equal(result, expected)
+
+    @given(before=hypothesis.strategies.integers(max_value=-1))
+    def test_error_before(self, *, before: int) -> None:
+        with raises(
+            _IsNearEventBeforeError, match=r"'Before' must be non-negative; got \-\d+"
+        ):
+            _ = is_near_event(before=before)
+
+    @given(after=hypothesis.strategies.integers(max_value=-1))
+    def test_error_after(self, *, after: int) -> None:
+        with raises(
+            _IsNearEventAfterError, match=r"'After' must be non-negative; got \-\d+"
+        ):
+            _ = is_near_event(after=after)
+
+
 class TestIsNullAndIsNotNullStructSeries:
     @given(
         case=sampled_from([
@@ -1877,6 +1981,98 @@ class TestNormal:
         assert series.is_finite().all()
 
 
+class TestReifyExprs:
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_one_expr(self, *, length: int, name: str) -> None:
+        expr = int_range(end=length).alias(name)
+        result = reify_exprs(expr)
+        assert isinstance(result, Expr)
+        result2 = (
+            int_range(end=length, eager=True)
+            .alias(f"_{name}")
+            .to_frame()
+            .with_columns(result)[name]
+        )
+        expected = int_range(end=length, eager=True).alias(name)
+        assert_series_equal(result2, expected)
+
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_one_series(self, *, length: int, name: str) -> None:
+        series = int_range(end=length, eager=True).alias(name)
+        result = reify_exprs(series)
+        assert isinstance(result, Series)
+        assert_series_equal(result, series)
+
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        names=pairs(text_ascii(), unique=True),
+    )
+    def test_two_exprs(self, *, length: int, names: tuple[str, str]) -> None:
+        name1, name2 = names
+        expr1 = int_range(end=length).alias(name1)
+        expr2 = int_range(end=length).alias(name2)
+        result = reify_exprs(expr1, expr2)
+        assert isinstance(result, Expr)
+        result2 = (
+            int_range(end=length, eager=True)
+            .alias(f"_{names}")
+            .to_frame()
+            .with_columns(result)[name1]
+        )
+        assert result2.name == name1
+        assert result2.dtype == Struct(dict.fromkeys(names, Int64))
+
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        names=pairs(text_ascii(), unique=True),
+    )
+    def test_one_expr_and_one_series(
+        self, *, length: int, names: tuple[str, str]
+    ) -> None:
+        name1, name2 = names
+        expr = int_range(end=length).alias(name1)
+        series = int_range(end=length, eager=True).alias(name2)
+        result = reify_exprs(expr, series)
+        assert isinstance(result, DataFrame)
+        assert result.schema == dict.fromkeys(names, Int64)
+
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        names=pairs(text_ascii(), unique=True),
+    )
+    def test_two_series(self, *, length: int, names: tuple[str, str]) -> None:
+        name1, name2 = names
+        series1 = int_range(end=length, eager=True).alias(name1)
+        series2 = int_range(end=length, eager=True).alias(name2)
+        result = reify_exprs(series1, series2)
+        assert isinstance(result, DataFrame)
+        expected = concat_series(series1, series2)
+        assert_frame_equal(result, expected)
+
+    def test_error_empty(self) -> None:
+        with raises(
+            _ReifyExprsEmptyError, match="At least 1 Expression or Series must be given"
+        ):
+            _ = reify_exprs()
+
+    @given(
+        lengths=pairs(hypothesis.strategies.integers(0, 10), unique=True),
+        names=pairs(text_ascii(), unique=True),
+    )
+    def test_error_non_unique(
+        self, *, lengths: tuple[int, int], names: tuple[str, str]
+    ) -> None:
+        series1, series2 = [
+            int_range(end=length, eager=True).alias(name)
+            for length, name in zip(lengths, names, strict=True)
+        ]
+        with raises(
+            _ReifyExprsSeriesNonUniqueError,
+            match=r"Series must contain exactly one length; got \d+, \d+ and perhaps more",
+        ):
+            _ = reify_exprs(series1, series2)
+
+
 class TestReplaceTimeZone:
     def test_datetime(self) -> None:
         now_utc = get_now()
@@ -2062,6 +2258,177 @@ class TestStructFromDataClass:
             _ = struct_from_dataclass(Example)
 
 
+class TestTryReifyExpr:
+    # expr
+
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_flat_expr(self, *, length: int, name: str) -> None:
+        expr = int_range(end=length).alias(name)
+        result = try_reify_expr(expr)
+        assert isinstance(result, Expr)
+        result2 = (
+            int_range(end=length, eager=True)
+            .alias(f"_{name}")
+            .to_frame()
+            .with_columns(result)[name]
+        )
+        expected = int_range(end=length, eager=True).alias(name)
+        assert_series_equal(result2, expected)
+
+    @given(length=hypothesis.strategies.integers(0, 10), names=pairs(text_ascii()))
+    def test_flat_expr_and_expr(self, *, length: int, names: tuple[str, str]) -> None:
+        name1, name2 = names
+        expr1 = int_range(end=length).alias(name1)
+        expr2 = int_range(end=length).alias(name2)
+        result = try_reify_expr(expr1, expr2)
+        assert isinstance(result, Expr)
+        result2 = (
+            int_range(end=length, eager=True)
+            .alias(f"_{name1}")
+            .to_frame()
+            .with_columns(result)[name1]
+        )
+        expected = int_range(end=length, eager=True).alias(name1)
+        assert_series_equal(result2, expected)
+
+    @given(length=hypothesis.strategies.integers(0, 10), names=pairs(text_ascii()))
+    def test_flat_expr_and_series(self, *, length: int, names: tuple[str, str]) -> None:
+        name1, name2 = names
+        expr = int_range(end=length).alias(name1)
+        series = int_range(end=length, eager=True).alias(name2)
+        result = try_reify_expr(expr, series)
+        assert isinstance(result, Series)
+        assert_series_equal(result, series.alias(name1))
+
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_struct_expr(self, *, length: int, name: str) -> None:
+        expr = struct(int_range(end=length).alias(name)).alias(name)
+        result = try_reify_expr(expr)
+        assert isinstance(result, Expr)
+        result2 = (
+            int_range(end=length, eager=True)
+            .alias(f"_{name}")
+            .to_frame()
+            .with_columns(result)[name]
+        )
+        expected = (
+            int_range(end=length, eager=True)
+            .alias(name)
+            .to_frame()
+            .select(struct(name))[name]
+        )
+        assert_series_equal(result2, expected)
+
+    @given(length=hypothesis.strategies.integers(0, 10), names=pairs(text_ascii()))
+    def test_struct_expr_and_expr(self, *, length: int, names: tuple[str, str]) -> None:
+        name1, name2 = names
+        expr1 = struct(int_range(end=length).alias(name1)).alias(name1)
+        expr2 = int_range(end=length).alias(name2)
+        result = try_reify_expr(expr1, expr2)
+        assert isinstance(result, Expr)
+        result2 = (
+            int_range(end=length, eager=True)
+            .alias(f"_{name1}")
+            .to_frame()
+            .with_columns(result)[name1]
+        )
+        expected = (
+            int_range(end=length, eager=True)
+            .alias(name1)
+            .to_frame()
+            .select(struct(name1))[name1]
+        )
+        assert_series_equal(result2, expected)
+
+    @given(length=hypothesis.strategies.integers(0, 10), names=pairs(text_ascii()))
+    def test_struct_expr_and_series(
+        self, *, length: int, names: tuple[str, str]
+    ) -> None:
+        name1, name2 = names
+        expr = struct(int_range(end=length).alias(name1)).alias(name1)
+        series = int_range(end=length, eager=True).alias(name2)
+        result = try_reify_expr(expr, series)
+        assert isinstance(result, Series)
+        expected = series.alias(name1).to_frame().select(struct(name1))[name1]
+        assert_series_equal(result, expected)
+
+    # series
+
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_flat_series(self, *, length: int, name: str) -> None:
+        series = int_range(end=length, eager=True).alias(name)
+        result = try_reify_expr(series)
+        assert isinstance(result, Series)
+        assert_series_equal(result, series)
+
+    @given(length=hypothesis.strategies.integers(0, 10), names=pairs(text_ascii()))
+    def test_flat_series_and_expr(self, *, length: int, names: tuple[str, str]) -> None:
+        name1, name2 = names
+        series = int_range(end=length, eager=True).alias(name1)
+        expr = int_range(end=length).alias(name2)
+        result = try_reify_expr(series, expr)
+        assert isinstance(result, Series)
+        assert_series_equal(result, series)
+
+    @given(length=hypothesis.strategies.integers(0, 10), names=pairs(text_ascii()))
+    def test_flat_series_and_series(
+        self, *, length: int, names: tuple[str, str]
+    ) -> None:
+        name1, name2 = names
+        series1 = int_range(end=length, eager=True).alias(name1)
+        series2 = int_range(end=length, eager=True).alias(name2)
+        result = try_reify_expr(series1, series2)
+        assert isinstance(result, Series)
+        assert_series_equal(result, series1)
+
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_struct_series(self, *, length: int, name: str) -> None:
+        series = (
+            int_range(end=length, eager=True)
+            .alias(name)
+            .to_frame()
+            .select(struct(name))[name]
+        )
+        assert isinstance(series.dtype, Struct)
+        result = try_reify_expr(series)
+        assert isinstance(result, Series)
+        assert_series_equal(result, series)
+
+    @given(length=hypothesis.strategies.integers(0, 10), names=pairs(text_ascii()))
+    def test_struct_series_and_expr(
+        self, *, length: int, names: tuple[str, str]
+    ) -> None:
+        name1, name2 = names
+        series = (
+            int_range(end=length, eager=True)
+            .alias(name1)
+            .to_frame()
+            .select(struct(name1))[name1]
+        )
+        assert isinstance(series.dtype, Struct)
+        expr = int_range(end=length).alias(name2)
+        result = try_reify_expr(series, expr)
+        assert isinstance(result, Series)
+        assert_series_equal(result, series)
+
+    @given(length=hypothesis.strategies.integers(0, 10), names=pairs(text_ascii()))
+    def test_struct_series_and_series(
+        self, *, length: int, names: tuple[str, str]
+    ) -> None:
+        name1, name2 = names
+        series1 = (
+            int_range(end=length, eager=True)
+            .alias(name1)
+            .to_frame()
+            .select(struct(name1))[name1]
+        )
+        assert isinstance(series1.dtype, Struct)
+        series2 = int_range(end=length).alias(name2)
+        result = try_reify_expr(series1, series2)
+        assert isinstance(result, Series)
+        assert_series_equal(result, series1)
+
+
 class TestUniform:
     @given(
         length=hypothesis.strategies.integers(0, 10),

{dycw_utilities-0.112.10 → dycw_utilities-0.112.11}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.112.10"
+__version__ = "0.112.11"

{dycw_utilities-0.112.10 → dycw_utilities-0.112.11}/src/utilities/polars.py

@@ -7,7 +7,7 @@ from collections.abc import Set as AbstractSet
 from contextlib import suppress
 from dataclasses import asdict, dataclass
 from functools import partial, reduce
-from itertools import chain
+from itertools import chain, product
 from math import ceil, log
 from pathlib import Path
 from typing import (
@@ -41,6 +41,7 @@ from polars import (
     Struct,
     UInt32,
     all_horizontal,
+    any_horizontal,
     col,
     concat,
     int_range,
@@ -1265,13 +1266,7 @@ def finite_ewm_mean(
     column = ensure_expr_or_series(column)
     mean = column.fill_null(value=0.0).rolling_mean(len(weights), weights=list(weights))
     expr = when(column.is_not_null()).then(mean)
-    match column:
-        case Expr():
-            return expr
-        case Series() as series:
-            return series.to_frame().with_columns(expr.alias(series.name))[series.name]
-        case _ as never:
-            assert_never(never)
+    return try_reify_expr(expr, column)
 
 
 @dataclass(kw_only=True)
@@ -1599,6 +1594,69 @@ def integers(
 ##
 
 
+@overload
+def is_near_event(
+    *exprs: ExprLike, before: int = 0, after: int = 0, **named_exprs: ExprLike
+) -> Expr: ...
+@overload
+def is_near_event(
+    *exprs: Series, before: int = 0, after: int = 0, **named_exprs: Series
+) -> Series: ...
+@overload
+def is_near_event(
+    *exprs: IntoExprColumn,
+    before: int = 0,
+    after: int = 0,
+    **named_exprs: IntoExprColumn,
+) -> Expr | Series: ...
+def is_near_event(
+    *exprs: IntoExprColumn,
+    before: int = 0,
+    after: int = 0,
+    **named_exprs: IntoExprColumn,
+) -> Expr | Series:
+    """Compute the rows near any event."""
+    if before <= -1:
+        raise _IsNearEventBeforeError(before=before)
+    if after <= -1:
+        raise _IsNearEventAfterError(after=after)
+    all_exprs = ensure_expr_or_series_many(*exprs, **named_exprs)
+    shifts = range(-before, after + 1)
+    if len(all_exprs) == 0:
+        near = lit(value=False, dtype=Boolean)
+    else:
+        near_exprs = (
+            e.shift(s).fill_null(value=False) for e, s in product(all_exprs, shifts)
+        )
+        near = any_horizontal(*near_exprs)
+    return try_reify_expr(near, *exprs, **named_exprs)
+
+
+@dataclass(kw_only=True, slots=True)
+class IsNearEventError(Exception): ...
+
+
+@dataclass(kw_only=True, slots=True)
+class _IsNearEventBeforeError(IsNearEventError):
+    before: int
+
+    @override
+    def __str__(self) -> str:
+        return f"'Before' must be non-negative; got {self.before}"
+
+
+@dataclass(kw_only=True, slots=True)
+class _IsNearEventAfterError(IsNearEventError):
+    after: int
+
+    @override
+    def __str__(self) -> str:
+        return f"'After' must be non-negative; got {self.after}"
+
+
+##
+
+
 def is_not_null_struct_series(series: Series, /) -> Series:
     """Check if a struct-dtype Series is not null as per the <= 1.1 definition."""
     try:
@@ -1791,6 +1849,71 @@ def normal(
 ##
 
 
+def reify_exprs(
+    *exprs: IntoExprColumn, **named_exprs: IntoExprColumn
+) -> Expr | Series | DataFrame:
+    """Reify a set of expressions."""
+    all_exprs = ensure_expr_or_series_many(*exprs, **named_exprs)
+    if len(all_exprs) == 0:
+        raise _ReifyExprsEmptyError from None
+    series = [s for s in all_exprs if isinstance(s, Series)]
+    lengths = {s.len() for s in series}
+    try:
+        length = one(lengths)
+    except OneEmptyError:
+        match len(all_exprs):
+            case 0:
+                raise ImpossibleCaseError(
+                    case=[f"{all_exprs=}"]
+                ) from None  # pragma: no cover
+            case 1:
+                return one(all_exprs)
+            case _:
+                return struct(*all_exprs)
+    except OneNonUniqueError as error:
+        raise _ReifyExprsSeriesNonUniqueError(
+            first=error.first, second=error.second
+        ) from None
+    df = (
+        int_range(end=length, eager=True)
+        .alias("_index")
+        .to_frame()
+        .with_columns(*all_exprs)
+        .drop("_index")
+    )
+    match len(df.columns):
+        case 0:
+            raise ImpossibleCaseError(case=[f"{df.columns=}"])  # pragma: no cover
+        case 1:
+            return df[one(df.columns)]
+        case _:
+            return df
+
+
+@dataclass(kw_only=True, slots=True)
+class ReifyExprsError(Exception): ...
+
+
+@dataclass(kw_only=True, slots=True)
+class _ReifyExprsEmptyError(ReifyExprsError):
+    @override
+    def __str__(self) -> str:
+        return "At least 1 Expression or Series must be given"
+
+
+@dataclass
+class _ReifyExprsSeriesNonUniqueError(ReifyExprsError):
+    first: int
+    second: int
+
+    @override
+    def __str__(self) -> str:
+        return f"Series must contain exactly one length; got {self.first}, {self.second} and perhaps more"
+
+
+##
+
+
 @overload
 def replace_time_zone(
     obj: Series, /, *, time_zone: TimeZoneLike | None = UTC
@@ -1930,6 +2053,28 @@ class _StructFromDataClassTypeError(StructFromDataClassError):
 ##
 
 
+def try_reify_expr(
+    expr: IntoExprColumn, /, *exprs: IntoExprColumn, **named_exprs: IntoExprColumn
+) -> Expr | Series:
+    """Try reify an expression."""
+    expr = ensure_expr_or_series(expr)
+    all_exprs = ensure_expr_or_series_many(*exprs, **named_exprs)
+    all_exprs = [e.alias(f"_{i}") for i, e in enumerate(all_exprs)]
+    result = reify_exprs(expr, *all_exprs)
+    match result:
+        case Expr():
+            return expr
+        case Series() as series:
+            return series
+        case DataFrame() as df:
+            return df[get_expr_name(df, expr)]
+        case _ as never:
+            assert_never(never)
+
+
+##
+
+
 def uniform(
     obj: int | Series | DataFrame,
     /,
@@ -2123,6 +2268,7 @@ __all__ = [
     "InsertAfterError",
     "InsertBeforeError",
     "InsertBetweenError",
+    "IsNearEventError",
     "IsNullStructSeriesError",
     "SetFirstRowAsColumnsError",
     "StructFromDataClassError",
@@ -2157,6 +2303,7 @@ __all__ = [
     "insert_before",
     "insert_between",
     "integers",
+    "is_near_event",
     "is_not_null_struct_series",
     "is_null_struct_series",
     "join",
@@ -2169,6 +2316,7 @@ __all__ = [
     "struct_dtype",
     "struct_from_dataclass",
     "touch",
+    "try_reify_expr",
     "uniform",
     "unique_element",
     "yield_struct_series_dataclasses",