dycw-utilities 0.159.7__tar.gz → 0.160.0__tar.gz

{dycw_utilities-0.159.7 → dycw_utilities-0.160.0}/PKG-INFO

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

{dycw_utilities-0.159.7 → dycw_utilities-0.160.0}/pyproject.toml

@@ -227,7 +227,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.159.7"
+version = "0.160.0"
 
 [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.7"
+current_version = "0.160.0"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.159.7 → dycw_utilities-0.160.0}/src/tests/test_polars.py

@@ -98,6 +98,7 @@ from utilities.polars import (
     OneColumnEmptyError,
     OneColumnNonUniqueError,
     RoundToFloatError,
+    SelectExactError,
     SetFirstRowAsColumnsError,
     TimePeriodDType,
     _check_polars_dataframe_predicates,
@@ -179,7 +180,7 @@ from utilities.polars import (
     join_into_periods,
     map_over_columns,
     nan_sum_agg,
-    nan_sum_cols,
+    nan_sum_horizonal,
     normal,
     number_of_decimals,
     offset_datetime,
@@ -191,6 +192,7 @@ from utilities.polars import (
     reify_exprs,
     replace_time_zone,
     round_to_float,
+    select_exact,
     serialize_dataframe,
     serialize_series,
     set_first_row_as_columns,
@@ -1561,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
 
@@ -2186,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_horizonal(*series)
+        assert result.item() == expected
 
 
 class TestNormal:
@@ -2467,6 +2465,39 @@ class TestRoundToFloat:
             _ = round_to_float(cast("Any", x), cast("Any", y))
 
 
+class TestSelectExact:
+    df: ClassVar[DataFrame] = DataFrame(
+        data=[(True, False), (False, None), (None, True)],
+        schema={"x": Boolean, "y": Boolean},
+        orient="row",
+    )
+
+    def test_adding_and_reordering(self) -> None:
+        df = select_exact(self.df, "y", ~col.x.alias("z"), "x")
+        expected = DataFrame(
+            data=[(False, False, True), (None, True, False), (True, None, None)],
+            schema={"y": Boolean, "z": Boolean, "x": Boolean},
+            orient="row",
+        )
+        assert_frame_equal(df, expected)
+
+    def test_adding_and_dropping(self) -> None:
+        df = select_exact(self.df, "y", ~col.x.alias("z"), drop="x")
+        expected = DataFrame(
+            data=[(False, False), (None, True), (True, None)],
+            schema={"y": Boolean, "z": Boolean},
+            orient="row",
+        )
+        assert_frame_equal(df, expected)
+
+    def test_error(self) -> None:
+        with raises(
+            SelectExactError,
+            match=r"All columns must be selected; got \['y'\] remaining",
+        ):
+            _ = select_exact(self.df, ~col.x.alias("z"), "x")
+
+
 class TestSerializeAndDeserializeDataFrame:
     cases: ClassVar[list[tuple[PolarsDataType, SearchStrategy[Any]]]] = [
         (Boolean, booleans()),

{dycw_utilities-0.159.7 → dycw_utilities-0.160.0}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.159.7"
+__version__ = "0.160.0"

{dycw_utilities-0.159.7 → dycw_utilities-0.160.0}/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_horizonal(*columns: Series) -> Series: ...
+@overload
+def nan_sum_horizonal(*columns: IntoExprColumn) -> ExprOrSeries: ...
+def nan_sum_horizonal(*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)
 
 
 ##
@@ -2446,6 +2424,31 @@ class RoundToFloatError(Exception):
 ##
 
 
+def select_exact(
+    df: DataFrame, /, *columns: IntoExprColumn, drop: MaybeIterable[str] | None = None
+) -> DataFrame:
+    """Select an exact set of columns from a DataFrame."""
+    names = [get_expr_name(df, c) for c in columns]
+    drop = set() if drop is None else set(always_iterable(drop))
+    union = set(names) | drop
+    extra = [c for c in df.columns if c not in union]
+    if len(extra) >= 1:
+        raise SelectExactError(columns=extra)
+    return df.select(*columns)
+
+
+@dataclass(kw_only=True, slots=True)
+class SelectExactError(Exception):
+    columns: list[str]
+
+    @override
+    def __str__(self) -> str:
+        return f"All columns must be selected; got {get_repr(self.columns)} remaining"
+
+
+##
+
+
 def set_first_row_as_columns(df: DataFrame, /) -> DataFrame:
     """Set the first row of a DataFrame as its columns."""
     try:
@@ -2654,6 +2657,7 @@ __all__ = [
     "OneColumnError",
     "OneColumnNonUniqueError",
     "RoundToFloatError",
+    "SelectExactError",
     "SetFirstRowAsColumnsError",
     "TimePeriodDType",
     "acf",
@@ -2698,7 +2702,7 @@ __all__ = [
     "join_into_periods",
     "map_over_columns",
     "nan_sum_agg",
-    "nan_sum_cols",
+    "nan_sum_horizonal",
     "normal",
     "number_of_decimals",
     "offset_datetime",
@@ -2709,6 +2713,7 @@ __all__ = [
     "read_series",
     "replace_time_zone",
     "round_to_float",
+    "select_exact",
     "serialize_dataframe",
     "set_first_row_as_columns",
     "struct_dtype",