dycw-utilities 0.153.2__py3-none-any.whl → 0.153.4__py3-none-any.whl

{dycw_utilities-0.153.2.dist-info → dycw_utilities-0.153.4.dist-info}/METADATA

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

{dycw_utilities-0.153.2.dist-info → dycw_utilities-0.153.4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=BOz9C-obZ9rIhyLSCVDsOVjnOeSdkl-YhkX0PSDH0-A,60
+utilities/__init__.py,sha256=BnvK6zqFzFO7qJ6fHdUA5-tJskYIpnc-PyYTQx-U3jo,60
 utilities/altair.py,sha256=92E2lCdyHY4Zb-vCw6rEJIsWdKipuu-Tu2ab1ufUfAk,9079
 utilities/asyncio.py,sha256=QXkTtugXkqtYt7Do23zgYErqzdp6jwzPpV_SP9fJ1gI,16780
 utilities/atomicwrites.py,sha256=tPo6r-Rypd9u99u66B9z86YBPpnLrlHtwox_8Z7T34Y,5790
@@ -45,7 +45,7 @@ utilities/parse.py,sha256=JcJn5yXKhIWXBCwgBdPsyu7Hvcuw6kyEdqvaebCaI9k,17951
 utilities/pathlib.py,sha256=qGuU8XPmdgGpy8tOMUgelfXx3kxI8h9IaV3TI_06QGE,8428
 utilities/pickle.py,sha256=MBT2xZCsv0pH868IXLGKnlcqNx2IRVKYNpRcqiQQqxw,653
 utilities/platform.py,sha256=pTn7gw6N4T6LdKrf0virwarof_mze9WtoQlrGMzhGVI,2798
-utilities/polars.py,sha256=uxV4liIQoCO92aERDO-AMfa8Es_-xCKgwZoyTO3PazQ,74410
+utilities/polars.py,sha256=jNUs038dJMZ0-r3rmHFClIzHlx3naA9klgYBvRDJi9w,77738
 utilities/polars_ols.py,sha256=Uc9V5kvlWZ5cU93lKZ-cfAKdVFFw81tqwLW9PxtUvMs,5618
 utilities/postgres.py,sha256=70BPcb_Na0LqDXJfqu0JwkHkGrwyuKpIVooSm6NlRn8,12467
 utilities/pottery.py,sha256=HJ96oLRarTP37Vhg0WTyB3yAu2hETeg6HgRmpDIqyUs,6581
@@ -87,8 +87,8 @@ utilities/zoneinfo.py,sha256=FBMcUQ4662Aq8SsuCL1OAhDQiyANmVjtb-C30DRrWoE,1966
 utilities/pytest_plugins/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 utilities/pytest_plugins/pytest_randomly.py,sha256=B1qYVlExGOxTywq2r1SMi5o7btHLk2PNdY_b1p98dkE,409
 utilities/pytest_plugins/pytest_regressions.py,sha256=9v8kAXDM2ycIXJBimoiF4EgrwbUvxTycFWJiGR_GHhM,1466
-dycw_utilities-0.153.2.dist-info/METADATA,sha256=ZoTBo07vfu7uVKDOazdKRsN6QUwefZDYSLO15V1AAD4,1696
-dycw_utilities-0.153.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.153.2.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
-dycw_utilities-0.153.2.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.153.2.dist-info/RECORD,,
+dycw_utilities-0.153.4.dist-info/METADATA,sha256=WLaEHDFvWCO7WFzRO4QPALUhD0yW9XUGz5Lx9x77BBM,1696
+dycw_utilities-0.153.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.153.4.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
+dycw_utilities-0.153.4.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.153.4.dist-info/RECORD,,

utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.153.2"
+__version__ = "0.153.4"

utilities/polars.py

@@ -33,6 +33,7 @@ from polars import (
     any_horizontal,
     col,
     concat,
+    datetime_range,
     int_range,
     lit,
     struct,
@@ -49,6 +50,7 @@ from polars.exceptions import (
 )
 from polars.schema import Schema
 from polars.testing import assert_frame_equal, assert_series_equal
+from whenever import ZonedDateTime
 
 from utilities.dataclasses import _YieldFieldsInstance, yield_fields
 from utilities.errors import ImpossibleCaseError
@@ -122,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")
@@ -129,6 +132,7 @@ DatetimeUSEastern = Datetime(time_zone="US/Eastern")
 DatetimeUTC = Datetime(time_zone="UTC")
 _FINITE_EWM_MIN_WEIGHT = 0.9999
 
+
 ##
 
 
@@ -238,11 +242,6 @@ def _acf_process_qstats_pvalues(qstats: NDArrayF, pvalues: NDArrayF, /) -> DataF
 ##
 
 
-# def acf_halflife(series: Series,/)
-
-##
-
-
 def adjust_frequencies(
     series: Series,
     /,
@@ -264,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}
@@ -427,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)
@@ -849,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)
 
@@ -874,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)
 
@@ -891,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:
@@ -963,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"],
@@ -975,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(
@@ -1180,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
 
@@ -1191,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())
@@ -1233,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,
     /,
@@ -1243,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(
@@ -1309,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)
@@ -1476,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:
@@ -1490,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:
@@ -1650,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)
@@ -2021,6 +2051,19 @@ def normal(
 ##
 
 
+def offset_datetime(
+    datetime: ZonedDateTime, offset: str, /, *, n: int = 1
+) -> ZonedDateTime:
+    """Offset a datetime as `polars` would."""
+    sr = Series(values=[datetime.py_datetime()])
+    for _ in range(n):
+        sr = sr.dt.offset_by(offset)
+    return ZonedDateTime.from_py_datetime(sr.item())
+
+
+##
+
+
 @overload
 def order_of_magnitude(column: ExprLike, /, *, round_: bool = False) -> Expr: ...
 @overload
@@ -2028,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()
@@ -2041,6 +2084,75 @@ def order_of_magnitude(
 ##
 
 
+@overload
+def period_range(
+    start: ZonedDateTime,
+    end_or_length: ZonedDateTime | int,
+    /,
+    *,
+    interval: str = "1d",
+    time_unit: TimeUnit | None = None,
+    time_zone: TimeZoneLike | None = None,
+    eager: Literal[True],
+) -> Series: ...
+@overload
+def period_range(
+    start: ZonedDateTime,
+    end_or_length: ZonedDateTime | int,
+    /,
+    *,
+    interval: str = "1d",
+    time_unit: TimeUnit | None = None,
+    time_zone: TimeZoneLike | None = None,
+    eager: Literal[False] = False,
+) -> Expr: ...
+@overload
+def period_range(
+    start: ZonedDateTime,
+    end_or_length: ZonedDateTime | int,
+    /,
+    *,
+    interval: str = "1d",
+    time_unit: TimeUnit | None = None,
+    time_zone: TimeZoneLike | None = None,
+    eager: bool = False,
+) -> Series | Expr: ...
+def period_range(
+    start: ZonedDateTime,
+    end_or_length: ZonedDateTime | int,
+    /,
+    *,
+    interval: str = "1d",
+    time_unit: TimeUnit | None = None,
+    time_zone: TimeZoneLike | None = None,
+    eager: bool = False,
+) -> Series | Expr:
+    """Construct a period range."""
+    time_zone_use = None if time_zone is None else ensure_time_zone(time_zone).key
+    match end_or_length:
+        case ZonedDateTime() as end:
+            ...
+        case int() as length:
+            end = offset_datetime(start, interval, n=length)
+        case never:
+            assert_never(never)
+    starts = datetime_range(
+        start.py_datetime(),
+        end.py_datetime(),
+        interval,
+        closed="left",
+        time_unit=time_unit,
+        time_zone=time_zone_use,
+        eager=eager,
+    ).alias("start")
+    ends = (starts.dt.offset_by(interval)).alias("end")
+    period = struct(starts, ends)
+    return try_reify_expr(period, starts, ends)
+
+
+##
+
+
 def reify_exprs(
     *exprs: IntoExprColumn, **named_exprs: IntoExprColumn
 ) -> Expr | Series | DataFrame:
@@ -2363,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)
@@ -2430,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")
@@ -2477,6 +2589,7 @@ __all__ = [
     "DatetimeUSEastern",
     "DatetimeUTC",
     "DropNullStructSeriesError",
+    "ExprOrSeries",
     "FiniteEWMMeanError",
     "GetDataTypeOrSeriesTimeZoneError",
     "GetSeriesNumberOfDecimalsError",
@@ -2489,6 +2602,10 @@ __all__ = [
     "StructFromDataClassError",
     "acf",
     "adjust_frequencies",
+    "all_dataframe_columns",
+    "all_series",
+    "any_dataframe_columns",
+    "any_series",
     "append_dataclass",
     "are_frames_equal",
     "bernoulli",
@@ -2529,7 +2646,9 @@ __all__ = [
     "nan_sum_agg",
     "nan_sum_cols",
     "normal",
+    "offset_datetime",
     "order_of_magnitude",
+    "period_range",
     "read_dataframe",
     "read_series",
     "replace_time_zone",