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

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

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

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

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

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

@@ -98,6 +98,7 @@ from utilities.polars import (
     OneColumnEmptyError,
     OneColumnNonUniqueError,
     RoundToFloatError,
+    SelectExactError,
     SetFirstRowAsColumnsError,
     TimePeriodDType,
     _check_polars_dataframe_predicates,
@@ -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,
@@ -2467,6 +2469,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.159.8}/src/utilities/__init__.py

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

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

@@ -2446,6 +2446,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 +2679,7 @@ __all__ = [
     "OneColumnError",
     "OneColumnNonUniqueError",
     "RoundToFloatError",
+    "SelectExactError",
     "SetFirstRowAsColumnsError",
     "TimePeriodDType",
     "acf",
@@ -2709,6 +2735,7 @@ __all__ = [
     "read_series",
     "replace_time_zone",
     "round_to_float",
+    "select_exact",
     "serialize_dataframe",
     "set_first_row_as_columns",
     "struct_dtype",