dycw-utilities 0.153.3__tar.gz → 0.153.5__tar.gz

{dycw_utilities-0.153.3 → dycw_utilities-0.153.5}/PKG-INFO

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

{dycw_utilities-0.153.3 → dycw_utilities-0.153.5}/pyproject.toml

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

{dycw_utilities-0.153.3 → dycw_utilities-0.153.5}/src/tests/test_polars.py

@@ -84,6 +84,7 @@ from utilities.polars import (
     DatetimeUSEastern,
     DatetimeUTC,
     DropNullStructSeriesError,
+    ExprOrSeries,
     FiniteEWMMeanError,
     InsertAfterError,
     InsertBeforeError,
@@ -132,6 +133,10 @@ from utilities.polars import (
     ac_halflife,
     acf,
     adjust_frequencies,
+    all_dataframe_columns,
+    all_series,
+    any_dataframe_columns,
+    any_series,
     append_dataclass,
     are_frames_equal,
     bernoulli,
@@ -205,6 +210,7 @@ if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Mapping, Sequence
     from zoneinfo import ZoneInfo
 
+    from _pytest.mark import ParameterSet
     from polars._typing import IntoExprColumn, PolarsDataType, SchemaDict
     from polars.datatypes import DataTypeClass
 
@@ -281,6 +287,56 @@ class TestAdjustFrequencies:
         assert isinstance(result, Series)
 
 
+class TestAnyAllDataFrameColumnsSeries:
+    cases: ClassVar[list[ParameterSet]] = [
+        param(int_range(end=pl.len()) % 2 == 0),
+        param(int_range(end=4, eager=True) % 2 == 0),
+    ]
+    series: ClassVar[Series] = Series(
+        name="x", values=[True, True, False, False], dtype=Boolean
+    )
+    df: ClassVar[DataFrame] = series.to_frame()
+    exp_all: ClassVar[Series] = Series(
+        name="x", values=[True, False, False, False], dtype=Boolean
+    )
+    exp_any: ClassVar[Series] = Series(
+        name="x", values=[True, True, True, False], dtype=Boolean
+    )
+    exp_empty: ClassVar[Series] = Series(
+        name="x", values=[True, False, True, False], dtype=Boolean
+    )
+
+    @mark.parametrize("column", cases)
+    def test_df_all(self, *, column: ExprOrSeries) -> None:
+        result = all_dataframe_columns(self.df, "x", column)
+        assert_series_equal(result, self.exp_all)
+
+    @mark.parametrize("column", cases)
+    def test_df_any(self, *, column: ExprOrSeries) -> None:
+        result = any_dataframe_columns(self.df, "x", column)
+        assert_series_equal(result, self.exp_any)
+
+    @mark.parametrize("column", cases)
+    def test_df_all_empty(self, *, column: ExprOrSeries) -> None:
+        result = all_dataframe_columns(self.df, column.alias("x"))
+        assert_series_equal(result, self.exp_empty)
+
+    @mark.parametrize("column", cases)
+    def test_df_any_empty(self, *, column: ExprOrSeries) -> None:
+        result = any_dataframe_columns(self.df, column.alias("x"))
+        assert_series_equal(result, self.exp_empty)
+
+    @mark.parametrize("column", cases)
+    def test_series_all(self, *, column: ExprOrSeries) -> None:
+        result = all_series(self.series, column)
+        assert_series_equal(result, self.exp_all)
+
+    @mark.parametrize("column", cases)
+    def test_series_any_any_any(self, *, column: ExprOrSeries) -> None:
+        result = any_series(self.series, column)
+        assert_series_equal(result, self.exp_any)
+
+
 class TestAppendDataClass:
     @given(
         data=fixed_dictionaries({

{dycw_utilities-0.153.3 → dycw_utilities-0.153.5}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.153.3"
+__version__ = "0.153.5"

{dycw_utilities-0.153.3 → dycw_utilities-0.153.5}/src/utilities/polars.py

@@ -124,6 +124,7 @@ if TYPE_CHECKING:
 
 
 type ExprLike = MaybeStr[Expr]
+type ExprOrSeries = Expr | Series
 DatetimeHongKong = Datetime(time_zone="Asia/Hong_Kong")
 DatetimeTokyo = Datetime(time_zone="Asia/Tokyo")
 DatetimeUSCentral = Datetime(time_zone="US/Central")
@@ -241,11 +242,6 @@ def _acf_process_qstats_pvalues(qstats: NDArrayF, pvalues: NDArrayF, /) -> DataF
 ##
 
 
-# def acf_halflife(series: Series,/)
-
-##
-
-
 def adjust_frequencies(
     series: Series,
     /,
@@ -267,6 +263,37 @@ def adjust_frequencies(
 ##
 
 
+def all_dataframe_columns(
+    df: DataFrame, expr: IntoExprColumn, /, *exprs: IntoExprColumn
+) -> Series:
+    """Return a DataFrame column with `AND` applied to additional exprs/series."""
+    name = get_expr_name(df, expr)
+    return df.select(all_horizontal(expr, *exprs).alias(name))[name]
+
+
+def any_dataframe_columns(
+    df: DataFrame, expr: IntoExprColumn, /, *exprs: IntoExprColumn
+) -> Series:
+    """Return a DataFrame column with `OR` applied to additional exprs/series."""
+    name = get_expr_name(df, expr)
+    return df.select(any_horizontal(expr, *exprs).alias(name))[name]
+
+
+def all_series(series: Series, /, *columns: ExprOrSeries) -> Series:
+    """Return a Series with `AND` applied to additional exprs/series."""
+    return all_dataframe_columns(series.to_frame(), series.name, *columns)
+
+
+def any_series(series: Series, /, *columns: ExprOrSeries) -> Series:
+    """Return a Series with `OR` applied to additional exprs/series."""
+    df = series.to_frame()
+    name = series.name
+    return df.select(any_horizontal(name, *columns).alias(name))[name]
+
+
+##
+
+
 def append_dataclass(df: DataFrame, obj: Dataclass, /) -> DataFrame:
     """Append a dataclass object to a DataFrame."""
     non_null_fields = {k: v for k, v in asdict(obj).items() if v is not None}
@@ -430,8 +457,8 @@ def ceil_datetime(column: ExprLike, every: ExprLike, /) -> Expr: ...
 @overload
 def ceil_datetime(column: Series, every: ExprLike, /) -> Series: ...
 @overload
-def ceil_datetime(column: IntoExprColumn, every: ExprLike, /) -> Expr | Series: ...
-def ceil_datetime(column: IntoExprColumn, every: ExprLike, /) -> Expr | Series:
+def ceil_datetime(column: IntoExprColumn, every: ExprLike, /) -> ExprOrSeries: ...
+def ceil_datetime(column: IntoExprColumn, every: ExprLike, /) -> ExprOrSeries:
     """Compute the `ceil` of a datetime column."""
     column = ensure_expr_or_series(column)
     rounded = column.dt.round(every)
@@ -852,13 +879,13 @@ def cross(
     up_or_down: Literal["up", "down"],
     other: Number | IntoExprColumn,
     /,
-) -> Expr | Series: ...
+) -> ExprOrSeries: ...
 def cross(
     expr: IntoExprColumn,
     up_or_down: Literal["up", "down"],
     other: Number | IntoExprColumn,
     /,
-) -> Expr | Series:
+) -> ExprOrSeries:
     """Compute when a cross occurs."""
     return _cross_or_touch(expr, "cross", up_or_down, other)
 
@@ -877,13 +904,13 @@ def touch(
     up_or_down: Literal["up", "down"],
     other: Number | IntoExprColumn,
     /,
-) -> Expr | Series: ...
+) -> ExprOrSeries: ...
 def touch(
     expr: IntoExprColumn,
     up_or_down: Literal["up", "down"],
     other: Number | IntoExprColumn,
     /,
-) -> Expr | Series:
+) -> ExprOrSeries:
     """Compute when a touch occurs."""
     return _cross_or_touch(expr, "touch", up_or_down, other)
 
@@ -894,7 +921,7 @@ def _cross_or_touch(
     up_or_down: Literal["up", "down"],
     other: Number | IntoExprColumn,
     /,
-) -> Expr | Series:
+) -> ExprOrSeries:
     """Compute when a column crosses/touches a threshold."""
     expr = ensure_expr_or_series(expr)
     match other:
@@ -966,7 +993,7 @@ def cross_rolling_quantile(
     weights: list[float] | None = None,
     min_samples: int | None = None,
     center: bool = False,
-) -> Expr | Series: ...
+) -> ExprOrSeries: ...
 def cross_rolling_quantile(
     expr: IntoExprColumn,
     up_or_down: Literal["up", "down"],
@@ -978,7 +1005,7 @@ def cross_rolling_quantile(
     weights: list[float] | None = None,
     min_samples: int | None = None,
     center: bool = False,
-) -> Expr | Series:
+) -> ExprOrSeries:
     """Compute when a column crosses its rolling quantile."""
     expr = ensure_expr_or_series(expr)
     rolling = expr.rolling_quantile(
@@ -1183,8 +1210,8 @@ def ensure_expr_or_series(column: ExprLike, /) -> Expr: ...
 @overload
 def ensure_expr_or_series(column: Series, /) -> Series: ...
 @overload
-def ensure_expr_or_series(column: IntoExprColumn, /) -> Expr | Series: ...
-def ensure_expr_or_series(column: IntoExprColumn, /) -> Expr | Series:
+def ensure_expr_or_series(column: IntoExprColumn, /) -> ExprOrSeries: ...
+def ensure_expr_or_series(column: IntoExprColumn, /) -> ExprOrSeries:
     """Ensure a column expression or Series is returned."""
     return col(column) if isinstance(column, str) else column
 
@@ -1194,7 +1221,7 @@ def ensure_expr_or_series(column: IntoExprColumn, /) -> Expr | Series:
 
 def ensure_expr_or_series_many(
     *columns: IntoExprColumn, **named_columns: IntoExprColumn
-) -> Sequence[Expr | Series]:
+) -> Sequence[ExprOrSeries]:
     """Ensure a set of column expressions and/or Series are returned."""
     args = map(ensure_expr_or_series, columns)
     kwargs = (ensure_expr_or_series(v).alias(k) for k, v in named_columns.items())
@@ -1236,7 +1263,7 @@ def finite_ewm_mean(
     half_life: float | None = None,
     alpha: float | None = None,
     min_weight: float = _FINITE_EWM_MIN_WEIGHT,
-) -> Expr | Series: ...
+) -> ExprOrSeries: ...
 def finite_ewm_mean(
     column: IntoExprColumn,
     /,
@@ -1246,7 +1273,7 @@ def finite_ewm_mean(
     half_life: float | None = None,
     alpha: float | None = None,
     min_weight: float = _FINITE_EWM_MIN_WEIGHT,
-) -> Expr | Series:
+) -> ExprOrSeries:
     """Compute a finite EWMA."""
     try:
         weights = _finite_ewm_weights(
@@ -1312,8 +1339,8 @@ def floor_datetime(column: ExprLike, every: ExprLike, /) -> Expr: ...
 @overload
 def floor_datetime(column: Series, every: ExprLike, /) -> Series: ...
 @overload
-def floor_datetime(column: IntoExprColumn, every: ExprLike, /) -> Expr | Series: ...
-def floor_datetime(column: IntoExprColumn, every: ExprLike, /) -> Expr | Series:
+def floor_datetime(column: IntoExprColumn, every: ExprLike, /) -> ExprOrSeries: ...
+def floor_datetime(column: IntoExprColumn, every: ExprLike, /) -> ExprOrSeries:
     """Compute the `floor` of a datetime column."""
     column = ensure_expr_or_series(column)
     rounded = column.dt.round(every)
@@ -1479,8 +1506,8 @@ def increasing_horizontal(*columns: ExprLike) -> Expr: ...
 @overload
 def increasing_horizontal(*columns: Series) -> Series: ...
 @overload
-def increasing_horizontal(*columns: IntoExprColumn) -> Expr | Series: ...
-def increasing_horizontal(*columns: IntoExprColumn) -> Expr | Series:
+def increasing_horizontal(*columns: IntoExprColumn) -> ExprOrSeries: ...
+def increasing_horizontal(*columns: IntoExprColumn) -> ExprOrSeries:
     """Check if a set of columns are increasing."""
     columns2 = ensure_expr_or_series_many(*columns)
     if len(columns2) == 0:
@@ -1493,8 +1520,8 @@ def decreasing_horizontal(*columns: ExprLike) -> Expr: ...
 @overload
 def decreasing_horizontal(*columns: Series) -> Series: ...
 @overload
-def decreasing_horizontal(*columns: IntoExprColumn) -> Expr | Series: ...
-def decreasing_horizontal(*columns: IntoExprColumn) -> Expr | Series:
+def decreasing_horizontal(*columns: IntoExprColumn) -> ExprOrSeries: ...
+def decreasing_horizontal(*columns: IntoExprColumn) -> ExprOrSeries:
     """Check if a set of columns are decreasing."""
     columns2 = ensure_expr_or_series_many(*columns)
     if len(columns2) == 0:
@@ -1653,13 +1680,13 @@ def is_near_event(
     before: int = 0,
     after: int = 0,
     **named_exprs: IntoExprColumn,
-) -> Expr | Series: ...
+) -> ExprOrSeries: ...
 def is_near_event(
     *exprs: IntoExprColumn,
     before: int = 0,
     after: int = 0,
     **named_exprs: IntoExprColumn,
-) -> Expr | Series:
+) -> ExprOrSeries:
     """Compute the rows near any event."""
     if before <= -1:
         raise _IsNearEventBeforeError(before=before)
@@ -2044,10 +2071,10 @@ def order_of_magnitude(column: Series, /, *, round_: bool = False) -> Series: ..
 @overload
 def order_of_magnitude(
     column: IntoExprColumn, /, *, round_: bool = False
-) -> Expr | Series: ...
+) -> ExprOrSeries: ...
 def order_of_magnitude(
     column: IntoExprColumn, /, *, round_: bool = False
-) -> Expr | Series:
+) -> ExprOrSeries:
     """Compute the order of magnitude of a column."""
     column = ensure_expr_or_series(column)
     result = column.abs().log10()
@@ -2448,7 +2475,7 @@ class _StructFromDataClassTypeError(StructFromDataClassError):
 
 def try_reify_expr(
     expr: IntoExprColumn, /, *exprs: IntoExprColumn, **named_exprs: IntoExprColumn
-) -> Expr | Series:
+) -> ExprOrSeries:
     """Try reify an expression."""
     expr = ensure_expr_or_series(expr)
     all_exprs = ensure_expr_or_series_many(*exprs, **named_exprs)
@@ -2515,8 +2542,8 @@ def week_num(column: ExprLike, /, *, start: WeekDay = "mon") -> Expr: ...
 @overload
 def week_num(column: Series, /, *, start: WeekDay = "mon") -> Series: ...
 @overload
-def week_num(column: IntoExprColumn, /, *, start: WeekDay = "mon") -> Expr | Series: ...
-def week_num(column: IntoExprColumn, /, *, start: WeekDay = "mon") -> Expr | Series:
+def week_num(column: IntoExprColumn, /, *, start: WeekDay = "mon") -> ExprOrSeries: ...
+def week_num(column: IntoExprColumn, /, *, start: WeekDay = "mon") -> ExprOrSeries:
     """Compute the week number of a date column."""
     column = ensure_expr_or_series(column)
     epoch = column.dt.epoch(time_unit="d").alias("epoch")
@@ -2562,6 +2589,7 @@ __all__ = [
     "DatetimeUSEastern",
     "DatetimeUTC",
     "DropNullStructSeriesError",
+    "ExprOrSeries",
     "FiniteEWMMeanError",
     "GetDataTypeOrSeriesTimeZoneError",
     "GetSeriesNumberOfDecimalsError",
@@ -2574,6 +2602,10 @@ __all__ = [
     "StructFromDataClassError",
     "acf",
     "adjust_frequencies",
+    "all_dataframe_columns",
+    "all_series",
+    "any_dataframe_columns",
+    "any_series",
     "append_dataclass",
     "are_frames_equal",
     "bernoulli",

{dycw_utilities-0.153.3 → dycw_utilities-0.153.5}/src/utilities/types.py

@@ -118,6 +118,7 @@ IPv6AddressLike = MaybeStr[IPv6Address]
 type MaybeCollection[T] = T | Collection[T]
 type MaybeIterable[T] = T | Iterable[T]
 type MaybeList[T] = T | list[T]
+type MaybeSet[T] = T | set[T] | frozenset[T]
 type SequenceLT[T] = list[T] | tuple[T, ...]
 # iterables - dervied
 type MaybeSequence[T] = T | SequenceLT[T]
@@ -303,6 +304,7 @@ __all__ = [
     "MaybeList",
     "MaybeSequence",
     "MaybeSequenceStr",
+    "MaybeSet",
     "MaybeStr",
     "MaybeType",
     "MonthDayLike",