dycw-utilities 0.166.15__tar.gz → 0.166.16__tar.gz

{dycw_utilities-0.166.15 → dycw_utilities-0.166.16}/PKG-INFO

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

{dycw_utilities-0.166.15 → dycw_utilities-0.166.16}/pyproject.toml

@@ -246,7 +246,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.166.15"
+version = "0.166.16"
 
 [project.entry-points.pytest11]
 pytest-randomly = "utilities.pytest_plugins.pytest_randomly"
@@ -279,7 +279,7 @@ test = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.166.15"
+current_version = "0.166.16"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.166.15 → dycw_utilities-0.166.16}/src/tests/test_polars.py

@@ -6,7 +6,7 @@ import itertools
 import math
 from dataclasses import dataclass, field
 from enum import auto
-from itertools import chain, repeat
+from itertools import chain
 from math import isfinite, nan
 from random import Random
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, assert_never, cast
@@ -49,6 +49,7 @@ from polars import (
     col,
     concat,
     date_range,
+    datetime_range,
     int_range,
     lit,
     struct,
@@ -58,7 +59,14 @@ from polars.schema import Schema
 from polars.testing import assert_frame_equal, assert_series_equal
 from pytest import mark, param, raises
 from scipy.stats import norm
-from whenever import DateDelta, DateTimeDelta, PlainDateTime, TimeDelta, ZonedDateTime
+from whenever import (
+    Date,
+    DateDelta,
+    DateTimeDelta,
+    PlainDateTime,
+    TimeDelta,
+    ZonedDateTime,
+)
 
 import tests.test_math
 import utilities.polars
@@ -168,6 +176,9 @@ from utilities.polars import (
     ensure_expr_or_series,
     ensure_expr_or_series_many,
     expr_to_series,
+    false_like,
+    filter_date,
+    filter_time,
     finite_ewm_mean,
     first_true_horizontal,
     get_data_type_or_series_time_zone,
@@ -208,6 +219,7 @@ from utilities.polars import (
     to_not_true,
     to_true,
     touch,
+    true_like,
     try_reify_expr,
     uniform,
     unique_element,
@@ -1479,6 +1491,42 @@ class TestExprToSeries:
         assert_series_equal(series, expected)
 
 
+class TestComputeDateFilter:
+    def test_main(self) -> None:
+        series = datetime_range(
+            start=ZonedDateTime(2024, 1, 1, tz=UTC.key).py_datetime(),
+            end=ZonedDateTime(2024, 1, 4, 12, tz=UTC.key).py_datetime(),
+            interval="12h",
+            eager=True,
+        ).alias("datetime")
+        assert len(series) == 8
+        result = filter_date(series, include=[Date(2024, 1, 2), Date(2024, 1, 3)])
+        expected = Series(
+            name="datetime",
+            values=[False, False, True, True, True, True, False, False],
+            dtype=Boolean,
+        )
+        assert_series_equal(result, expected)
+
+
+class TestComputeTimeFilter:
+    def test_main(self) -> None:
+        series = datetime_range(
+            start=ZonedDateTime(2024, 1, 1, tz=UTC.key).py_datetime(),
+            end=ZonedDateTime(2024, 1, 3, 0, tz=UTC.key).py_datetime(),
+            interval="6h",
+            eager=True,
+        ).alias("datetime")
+        assert len(series) == 9
+        result = filter_time(series, include=[(whenever.Time(6), whenever.Time(12))])
+        expected = Series(
+            name="datetime",
+            values=[False, True, True, False, False, True, True, False, False],
+            dtype=Boolean,
+        )
+        assert_series_equal(result, expected)
+
+
 class TestFiniteEWMMean:
     alpha_0_75_values: ClassVar[list[float]] = [
         -8.269850726503885,
@@ -1838,7 +1886,9 @@ class TestIsNearEvent:
     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
+            name="z",
+            values=list(itertools.repeat(object=False, times=10)),
+            dtype=Boolean,
         )
         assert_series_equal(result, expected)
 
@@ -2783,6 +2833,60 @@ class TestToTrueAndFalse:
         assert_series_equal(result, exp_series)
 
 
+class TestTrueLikeAndFalseLike:
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_true_expr(self, *, length: int, name: str) -> None:
+        expr = int_range(end=length).alias(name)
+        result = true_like(expr)
+        assert isinstance(result, Expr)
+        result2 = (
+            int_range(end=length, eager=True)
+            .alias(f"_{name}")
+            .to_frame()
+            .with_columns(result)[name]
+        )
+        expected = pl.repeat(value=True, n=length, dtype=Boolean, eager=True).alias(
+            name
+        )
+        assert_series_equal(result2, expected)
+
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_true_series(self, *, length: int, name: str) -> None:
+        series = int_range(end=length, eager=True).alias(name)
+        result = true_like(series)
+        assert isinstance(result, Series)
+        expected = pl.repeat(value=True, n=length, dtype=Boolean, eager=True).alias(
+            name
+        )
+        assert_series_equal(result, expected)
+
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_false_expr(self, *, length: int, name: str) -> None:
+        expr = int_range(end=length).alias(name)
+        result = false_like(expr)
+        assert isinstance(result, Expr)
+        result2 = (
+            int_range(end=length, eager=True)
+            .alias(f"_{name}")
+            .to_frame()
+            .with_columns(result)[name]
+        )
+        expected = pl.repeat(value=False, n=length, dtype=Boolean, eager=True).alias(
+            name
+        )
+        assert_series_equal(result2, expected)
+
+    @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
+    def test_false_series(self, *, length: int, name: str) -> None:
+        series = int_range(end=length, eager=True).alias(name)
+        result = false_like(series)
+        assert isinstance(result, Series)
+        expected = pl.repeat(value=False, n=length, dtype=Boolean, eager=True).alias(
+            name
+        )
+        assert_series_equal(result, expected)
+
+
 class TestTryReifyExpr:
     # expr
 

{dycw_utilities-0.166.15 → dycw_utilities-0.166.16}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.166.15"
+__version__ = "0.166.16"

{dycw_utilities-0.166.15 → dycw_utilities-0.166.16}/src/utilities/polars.py

@@ -70,6 +70,7 @@ from utilities.iterables import (
     check_supermapping,
     is_iterable_not_str,
     one,
+    resolve_include_and_exclude,
 )
 from utilities.json import write_formatted_json
 from utilities.math import (
@@ -1269,6 +1270,111 @@ def expr_to_series(expr: Expr, /) -> Series:
 ##
 
 
+@overload
+def filter_date(
+    column: ExprLike = "datetime",
+    /,
+    *,
+    time_zone: ZoneInfo | None = None,
+    include: MaybeIterable[whenever.Date] | None = None,
+    exclude: MaybeIterable[whenever.Date] | None = None,
+) -> Expr: ...
+@overload
+def filter_date(
+    column: Series,
+    /,
+    *,
+    time_zone: ZoneInfo | None = None,
+    include: MaybeIterable[whenever.Date] | None = None,
+    exclude: MaybeIterable[whenever.Date] | None = None,
+) -> Series: ...
+@overload
+def filter_date(
+    column: IntoExprColumn = "datetime",
+    /,
+    *,
+    time_zone: ZoneInfo | None = None,
+    include: MaybeIterable[whenever.Date] | None = None,
+    exclude: MaybeIterable[whenever.Date] | None = None,
+) -> ExprOrSeries: ...
+def filter_date(
+    column: IntoExprColumn = "datetime",
+    /,
+    *,
+    time_zone: ZoneInfo | None = None,
+    include: MaybeIterable[whenever.Date] | None = None,
+    exclude: MaybeIterable[whenever.Date] | None = None,
+) -> ExprOrSeries:
+    """Compute the filter based on a set of dates."""
+    column = ensure_expr_or_series(column)
+    if time_zone is not None:
+        column = column.dt.convert_time_zone(time_zone.key)
+    keep = true_like(column)
+    date = column.dt.date()
+    include, exclude = resolve_include_and_exclude(include=include, exclude=exclude)
+    if include is not None:
+        keep &= date.is_in([d.py_date() for d in include])
+    if exclude is not None:
+        keep &= ~date.is_in([d.py_date() for d in exclude])
+    return try_reify_expr(keep, column)
+
+
+@overload
+def filter_time(
+    column: ExprLike = "datetime",
+    /,
+    *,
+    time_zone: ZoneInfo | None = None,
+    include: MaybeIterable[tuple[whenever.Time, whenever.Time]] | None = None,
+    exclude: MaybeIterable[tuple[whenever.Time, whenever.Time]] | None = None,
+) -> Expr: ...
+@overload
+def filter_time(
+    column: Series,
+    /,
+    *,
+    time_zone: ZoneInfo | None = None,
+    include: MaybeIterable[tuple[whenever.Time, whenever.Time]] | None = None,
+    exclude: MaybeIterable[tuple[whenever.Time, whenever.Time]] | None = None,
+) -> Series: ...
+@overload
+def filter_time(
+    column: IntoExprColumn = "datetime",
+    /,
+    *,
+    time_zone: ZoneInfo | None = None,
+    include: MaybeIterable[tuple[whenever.Time, whenever.Time]] | None = None,
+    exclude: MaybeIterable[tuple[whenever.Time, whenever.Time]] | None = None,
+) -> ExprOrSeries: ...
+def filter_time(
+    column: IntoExprColumn = "datetime",
+    /,
+    *,
+    time_zone: ZoneInfo | None = None,
+    include: MaybeIterable[tuple[whenever.Time, whenever.Time]] | None = None,
+    exclude: MaybeIterable[tuple[whenever.Time, whenever.Time]] | None = None,
+) -> ExprOrSeries:
+    """Compute the filter based on a set of times."""
+    column = ensure_expr_or_series(column)
+    if time_zone is not None:
+        column = column.dt.convert_time_zone(time_zone.key)
+    keep = true_like(column)
+    time = column.dt.time()
+    include, exclude = resolve_include_and_exclude(include=include, exclude=exclude)
+    if include is not None:
+        keep &= any_horizontal(
+            time.is_between(s.py_time(), e.py_time()) for s, e in include
+        )
+    if exclude is not None:
+        keep &= ~any_horizontal(
+            time.is_between(s.py_time(), e.py_time()) for s, e in exclude
+        )
+    return try_reify_expr(keep, column)
+
+
+##
+
+
 @overload
 def finite_ewm_mean(
     column: ExprLike,
@@ -2643,6 +2749,33 @@ def to_not_false(column: IntoExprColumn, /) -> ExprOrSeries:
 ##
 
 
+@overload
+def true_like(column: ExprLike, /) -> Expr: ...
+@overload
+def true_like(column: Series, /) -> Series: ...
+@overload
+def true_like(column: IntoExprColumn, /) -> ExprOrSeries: ...
+def true_like(column: IntoExprColumn, /) -> ExprOrSeries:
+    """Compute a column of `True` values."""
+    column = ensure_expr_or_series(column)
+    return column.is_null() | column.is_not_null()
+
+
+@overload
+def false_like(column: ExprLike, /) -> Expr: ...
+@overload
+def false_like(column: Series, /) -> Series: ...
+@overload
+def false_like(column: IntoExprColumn, /) -> ExprOrSeries: ...
+def false_like(column: IntoExprColumn, /) -> ExprOrSeries:
+    """Compute a column of `False` values."""
+    column = ensure_expr_or_series(column)
+    return column.is_null() & column.is_not_null()
+
+
+##
+
+
 def try_reify_expr(
     expr: IntoExprColumn, /, *exprs: IntoExprColumn, **named_exprs: IntoExprColumn
 ) -> ExprOrSeries:
@@ -2798,6 +2931,9 @@ __all__ = [
     "ensure_expr_or_series",
     "ensure_expr_or_series_many",
     "expr_to_series",
+    "false_like",
+    "filter_date",
+    "filter_time",
     "finite_ewm_mean",
     "first_true_horizontal",
     "get_data_type_or_series_time_zone",
@@ -2838,6 +2974,7 @@ __all__ = [
     "to_not_true",
     "to_true",
     "touch",
+    "true_like",
     "try_reify_expr",
     "uniform",
     "unique_element",