dycw-utilities 0.159.8__tar.gz → 0.160.1__tar.gz

{dycw_utilities-0.159.8 → dycw_utilities-0.160.1}/PKG-INFO

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

{dycw_utilities-0.159.8 → dycw_utilities-0.160.1}/pyproject.toml

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

{dycw_utilities-0.159.8 → dycw_utilities-0.160.1}/src/tests/test_polars.py

@@ -180,7 +180,7 @@ from utilities.polars import (
     join_into_periods,
     map_over_columns,
     nan_sum_agg,
-    nan_sum_cols,
+    nan_sum_horizontal,
     normal,
     number_of_decimals,
     offset_datetime,
@@ -1563,16 +1563,16 @@ class TestFiniteEWMWeights:
 
 class TestFirstTrueHorizontal:
     @mark.parametrize(
-        ("values", "expected"),
+        ("x", "y", "z", "expected"),
         [
-            param([True, True, True], 0),
-            param([False, True, True], 1),
-            param([False, False, True], 2),
-            param([False, False, False], None),
+            param(True, True, True, 0),
+            param(False, True, True, 1),
+            param(False, False, True, 2),
+            param(False, False, False, None),
         ],
     )
-    def test_main(self, *, values: list[bool], expected: int | None) -> None:
-        series = [Series(values=[v], dtype=Boolean) for v in values]
+    def test_main(self, *, x: bool, y: bool, z: bool, expected: int | None) -> None:
+        series = [Series(values=[i], dtype=Boolean) for i in [x, y, z]]
         result = first_true_horizontal(*series)
         assert result.item() == expected
 
@@ -2188,26 +2188,22 @@ class TestNanSumAgg:
         assert result["value"].item() == expected
 
 
-class TestNanSumCols:
-    @given(
-        case=sampled_from([(None, None, None), (None, 0, 0), (0, None, 0), (1, 2, 3)]),
-        x_kind=sampled_from(["str", "column"]),
-        y_kind=sampled_from(["str", "column"]),
+class TestNanSumHorizontal:
+    @mark.parametrize(
+        ("x", "y", "z", "expected"),
+        [
+            param(1, 2, 3, 6),
+            param(None, 2, 3, 5),
+            param(None, None, 3, 3),
+            param(None, None, None, None),
+        ],
     )
     def test_main(
-        self,
-        *,
-        case: tuple[int | None, int | None, int | None],
-        x_kind: Literal["str", "column"],
-        y_kind: Literal["str", "column"],
+        self, *, x: int | None, y: int | None, z: int | None, expected: int | None
     ) -> None:
-        x, y, expected = case
-        x_use = "x" if x_kind == "str" else col("x")
-        y_use = "y" if y_kind == "str" else col("y")
-        df = DataFrame(
-            data=[(x, y)], schema={"x": Int64, "y": Int64}, orient="row"
-        ).with_columns(z=nan_sum_cols(x_use, y_use))
-        assert df["z"].item() == expected
+        series = [Series(values=[i], dtype=Int64) for i in [x, y, z]]
+        result = nan_sum_horizontal(*series)
+        assert result.item() == expected
 
 
 class TestNormal:

{dycw_utilities-0.159.8 → dycw_utilities-0.160.1}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.159.8"
+__version__ = "0.160.1"

{dycw_utilities-0.159.8 → dycw_utilities-0.160.1}/src/utilities/polars.py

@@ -1325,18 +1325,6 @@ class _FiniteEWMWeightsError(Exception):
 ##
 
 
-@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
@@ -1931,27 +1919,17 @@ def nan_sum_agg(column: str | Expr, /, *, dtype: PolarsDataType | None = None) -
 ##
 
 
-def nan_sum_cols(
-    column: str | Expr, *columns: str | Expr, dtype: PolarsDataType | None = None
-) -> Expr:
+@overload
+def nan_sum_horizontal(*columns: Series) -> Series: ...
+@overload
+def nan_sum_horizontal(*columns: IntoExprColumn) -> ExprOrSeries: ...
+def nan_sum_horizontal(*columns: IntoExprColumn) -> ExprOrSeries:
     """Nan sum across columns."""
-    all_columns = chain([column], columns)
-    all_exprs = (
-        col(column) if isinstance(column, str) else column for column in all_columns
+    columns2 = ensure_expr_or_series_many(*columns)
+    expr = when(any_horizontal(*(c.is_not_null() for c in columns2))).then(
+        sum_horizontal(*columns2)
     )
-
-    def func(x: Expr, y: Expr, /) -> Expr:
-        return (
-            when(x.is_not_null() & y.is_not_null())
-            .then(x + y)
-            .when(x.is_not_null() & y.is_null())
-            .then(x)
-            .when(x.is_null() & y.is_not_null())
-            .then(y)
-            .otherwise(lit(None, dtype=dtype))
-        )
-
-    return reduce(func, all_exprs)
+    return try_reify_expr(expr, *columns2)
 
 
 ##
@@ -2724,7 +2702,7 @@ __all__ = [
     "join_into_periods",
     "map_over_columns",
     "nan_sum_agg",
-    "nan_sum_cols",
+    "nan_sum_horizontal",
     "normal",
     "number_of_decimals",
     "offset_datetime",