dycw-utilities 0.159.4__py3-none-any.whl → 0.159.5__py3-none-any.whl

{dycw_utilities-0.159.4.dist-info → dycw_utilities-0.159.5.dist-info}/METADATA

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

{dycw_utilities-0.159.4.dist-info → dycw_utilities-0.159.5.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=8PwC9PN_0tTaC-HGuAmY0sqVwC7bKnFBwSDdjEvRwnk,60
+utilities/__init__.py,sha256=y1aZeGKpHWH6gkJPUx6DJDCZYIPbZ7BynHdMXlgmmME,60
 utilities/altair.py,sha256=92E2lCdyHY4Zb-vCw6rEJIsWdKipuu-Tu2ab1ufUfAk,9079
 utilities/asyncio.py,sha256=PUedzQ5deqlSECQ33sam9cRzI9TnygHz3FdOqWJWPTM,15288
 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=4Qc9lv0SZEpjXFKANUQBb7nA55sC4zaHtlZEJXTBALg,76896
+utilities/polars.py,sha256=QTHk58M2dwIOCSR1JYttlrblUAw8Ihn7M8gAH4CtrGU,79542
 utilities/polars_ols.py,sha256=Uc9V5kvlWZ5cU93lKZ-cfAKdVFFw81tqwLW9PxtUvMs,5618
 utilities/postgres.py,sha256=ynCTTaF-bVEOSW-KEAR-dlLh_hYjeVVjm__-4pEU8Zk,12269
 utilities/pottery.py,sha256=ggMN72Y7wx7Js8VN6eyNyodpm8TIYqZHGghkDPXIVWk,3949
@@ -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.159.4.dist-info/METADATA,sha256=WLxPLavB-QLjzJhQr7T5-LbVJNRA_PlRiBtCJTT9FmQ,1643
-dycw_utilities-0.159.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.159.4.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
-dycw_utilities-0.159.4.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.159.4.dist-info/RECORD,,
+dycw_utilities-0.159.5.dist-info/METADATA,sha256=HkJFj45zxD0WbiOwR0QfrBq9RjjFHmI8aieoNk49Voo,1643
+dycw_utilities-0.159.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.159.5.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
+dycw_utilities-0.159.5.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.159.5.dist-info/RECORD,,

utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.159.4"
+__version__ = "0.159.5"

utilities/polars.py

@@ -57,6 +57,7 @@ import utilities.math
 from utilities.dataclasses import yield_fields
 from utilities.errors import ImpossibleCaseError
 from utilities.functions import (
+    get_class_name,
     is_dataclass_class,
     is_dataclass_instance,
     make_isinstance,
@@ -1324,10 +1325,27 @@ class _FiniteEWMWeightsError(Exception):
 ##
 
 
-def first_true_horizontal(df: DataFrame, /) -> Series:
+@overload
+def first_true_horizontal(column: Series, /) -> Series: ...
+@overload
+def first_true_horizontal(column1: Series, column2: Series, /) -> Series: ...
+@overload
+def first_true_horizontal(
+    column1: Series, column2: Series, column3: Series, /
+) -> Series: ...
+@overload
+def first_true_horizontal(
+    column1: Series, column2: Series, column3: Series, column4: Series, /
+) -> Series: ...
+@overload
+def first_true_horizontal(*columns: Series) -> Series: ...
+@overload
+def first_true_horizontal(*columns: IntoExprColumn) -> ExprOrSeries: ...
+def first_true_horizontal(*columns: IntoExprColumn) -> ExprOrSeries:
     """Get the index of the first true in each row."""
-    expr = when(any_horizontal(pl.all())).then(concat_list(pl.all()).list.arg_max())
-    return one_column(df.select(expr))
+    columns2 = ensure_expr_or_series_many(*columns)
+    expr = when(any_horizontal(*columns2)).then(concat_list(*columns2).list.arg_max())
+    return try_reify_expr(expr, *columns2)
 
 
 ##
@@ -1972,19 +1990,31 @@ def normal(
 ##
 
 
+@overload
+def number_of_decimals(
+    series: ExprLike, /, *, max_decimals: int = MAX_DECIMALS
+) -> Expr: ...
+@overload
 def number_of_decimals(
     series: Series, /, *, max_decimals: int = MAX_DECIMALS
-) -> Series:
+) -> Series: ...
+@overload
+def number_of_decimals(
+    series: IntoExprColumn, /, *, max_decimals: int = MAX_DECIMALS
+) -> ExprOrSeries: ...
+def number_of_decimals(
+    series: IntoExprColumn, /, *, max_decimals: int = MAX_DECIMALS
+) -> ExprOrSeries:
     """Get the number of decimals."""
+    series = ensure_expr_or_series(series)
     frac = series - series.floor()
-    results = [
+    results = (
         _number_of_decimals_check_scale(frac, s) for s in range(max_decimals + 1)
-    ]
-    df_results = concat_series(*results)
-    return first_true_horizontal(df_results)
+    )
+    return first_true_horizontal(*results)
 
 
-def _number_of_decimals_check_scale(frac: Series, scale: int, /) -> Series:
+def _number_of_decimals_check_scale(frac: ExprOrSeries, scale: int, /) -> ExprOrSeries:
     scaled = 10**scale * frac
     return is_close(scaled, scaled.round()).alias(str(scale))
 
@@ -2346,19 +2376,63 @@ def round_to_float(
 ) -> Expr: ...
 @overload
 def round_to_float(
-    x: Series, y: float, /, *, mode: RoundMode = "half_to_even"
+    x: Series, y: float | ExprOrSeries, /, *, mode: RoundMode = "half_to_even"
+) -> Series: ...
+@overload
+def round_to_float(
+    x: ExprLike, y: Series, /, *, mode: RoundMode = "half_to_even"
 ) -> Series: ...
 @overload
 def round_to_float(
-    x: IntoExprColumn, y: float, /, *, mode: RoundMode = "half_to_even"
+    x: ExprLike, y: Expr, /, *, mode: RoundMode = "half_to_even"
+) -> Expr: ...
+@overload
+def round_to_float(
+    x: IntoExprColumn, y: float | Series, /, *, mode: RoundMode = "half_to_even"
 ) -> ExprOrSeries: ...
 def round_to_float(
-    x: IntoExprColumn, y: float, /, *, mode: RoundMode = "half_to_even"
+    x: IntoExprColumn, y: float | IntoExprColumn, /, *, mode: RoundMode = "half_to_even"
 ) -> ExprOrSeries:
     """Round a column to the nearest multiple of another float."""
     x = ensure_expr_or_series(x)
-    z = (x / y).round(mode=mode) * y
-    return z.round(decimals=utilities.math.number_of_decimals(y) + 1)
+    y = y if isinstance(y, int | float) else ensure_expr_or_series(y)
+    match x, y:
+        case Expr() | Series(), int() | float():
+            z = (x / y).round(mode=mode) * y
+            return z.round(decimals=utilities.math.number_of_decimals(y) + 1)
+        case Series(), Expr() | Series():
+            df = (
+                x.to_frame()
+                .with_columns(y)
+                .with_columns(number_of_decimals(y).alias("_decimals"))
+                .with_row_index(name="_index")
+                .group_by("_decimals")
+                .map_groups(_round_to_float_one)
+                .sort("_index")
+            )
+            return df[df.columns[1]]
+        case Expr(), Series():
+            df = y.to_frame().with_columns(x)
+            return round_to_float(df[df.columns[1]], df[df.columns[0]], mode=mode)
+        case Expr(), Expr() | str():
+            raise RoundToFloatError(x=x, y=y)
+        case never:
+            assert_never(never)
+
+
+def _round_to_float_one(df: DataFrame, /) -> DataFrame:
+    decimals: int = df["_decimals"].unique().item()
+    return df.with_columns(col(df.columns[1]).round(decimals=decimals))
+
+
+@dataclass(kw_only=True, slots=True)
+class RoundToFloatError(Exception):
+    x: IntoExprColumn
+    y: IntoExprColumn
+
+    @override
+    def __str__(self) -> str:
+        return f"At least 1 of the dividend and/or divisor must be a Series; got {get_class_name(self.x)!r} and {get_class_name(self.y)!r}"
 
 
 ##
@@ -2571,6 +2645,7 @@ __all__ = [
     "OneColumnEmptyError",
     "OneColumnError",
     "OneColumnNonUniqueError",
+    "RoundToFloatError",
     "SetFirstRowAsColumnsError",
     "TimePeriodDType",
     "acf",