dycw-utilities 0.165.2__tar.gz → 0.166.0__tar.gz

{dycw_utilities-0.165.2 → dycw_utilities-0.166.0}/PKG-INFO

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

{dycw_utilities-0.165.2 → dycw_utilities-0.166.0}/pyproject.toml

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

{dycw_utilities-0.165.2 → dycw_utilities-0.166.0}/src/tests/test_polars.py

@@ -83,7 +83,6 @@ from utilities.hypothesis import (
 from utilities.numpy import DEFAULT_RNG
 from utilities.pathlib import PWD
 from utilities.polars import (
-    AppendDataClassError,
     BooleanValueCountsError,
     ColumnsToDictError,
     DatePeriodDType,
@@ -102,6 +101,10 @@ from utilities.polars import (
     SelectExactError,
     SetFirstRowAsColumnsError,
     TimePeriodDType,
+    _AppendRowExtraKeysError,
+    _AppendRowMissingKeysError,
+    _AppendRowNullColumnsError,
+    _AppendRowPredicateError,
     _check_polars_dataframe_predicates,
     _check_polars_dataframe_schema_list,
     _check_polars_dataframe_schema_set,
@@ -145,7 +148,7 @@ from utilities.polars import (
     all_series,
     any_dataframe_columns,
     any_series,
-    append_dataclass,
+    append_row,
     are_frames_equal,
     bernoulli,
     boolean_value_counts,
@@ -194,6 +197,7 @@ from utilities.polars import (
     reify_exprs,
     replace_time_zone,
     round_to_float,
+    search_period,
     select_exact,
     serialize_dataframe,
     serialize_series,
@@ -224,6 +228,7 @@ from utilities.whenever import (
     get_now,
     get_now_plain,
     get_today,
+    to_zoned_date_time,
 )
 from utilities.zoneinfo import UTC, to_time_zone_name
 
@@ -358,91 +363,89 @@ class TestAnyAllDataFrameColumnsSeries:
         assert_series_equal(result, self.exp_any)
 
 
-class TestAppendDataClass:
-    @given(
-        data=fixed_dictionaries({
-            "a": int64s() | none(),
-            "b": floats() | none(),
-            "c": text_ascii() | none(),
-        })
-    )
-    def test_columns_and_fields_equal(self, *, data: StrMapping) -> None:
-        df = DataFrame(schema={"a": Int64, "b": Float64, "c": String})
-
-        @dataclass(kw_only=True, slots=True)
-        class Row:
-            a: int | None = None
-            b: float | None = None
-            c: str | None = None
-
-        row = Row(**data)
-        result = append_dataclass(df, row)
-        height = 0 if (row.a is None) and (row.b is None) and (row.c is None) else 1
-        check_polars_dataframe(result, height=height, schema_list=df.schema)
-
-    @given(data=fixed_dictionaries({"a": int64s() | none(), "b": floats() | none()}))
-    def test_extra_column(self, *, data: StrMapping) -> None:
-        df = DataFrame(schema={"a": Int64, "b": Float64, "c": String})
+class TestAppendRow:
+    rows: ClassVar[list[tuple[int, int]]] = [(i, 2 * i) for i in range(3)]
+    schema: ClassVar[SchemaDict] = {"x": Int64, "y": Int64}
+    df: ClassVar[DataFrame] = DataFrame(data=rows, schema=schema, orient="row")
 
-        @dataclass(kw_only=True, slots=True)
-        class Row:
-            a: int | None = None
-            b: float | None = None
+    @mark.parametrize(
+        "row", [param({"x": 3, "y": 6}), param({"x": 3, "y": 6, "z": None})]
+    )
+    def test_main(self, *, row: StrMapping) -> None:
+        result = append_row(self.df, row)
+        expected = DataFrame(
+            data=[*self.rows, (3, 6)], schema=self.schema, orient="row"
+        )
+        assert_frame_equal(result, expected)
 
-        row = Row(**data)
-        result = append_dataclass(df, row)
-        height = 0 if (row.a is None) and (row.b is None) else 1
-        check_polars_dataframe(result, height=height, schema_list=df.schema)
+    def test_missing_key(self) -> None:
+        row = {"x": 3}
+        result = append_row(self.df, row)
+        expected = DataFrame(
+            data=[*self.rows, (3, None)], schema=self.schema, orient="row"
+        )
+        assert_frame_equal(result, expected)
 
-    @given(data=fixed_dictionaries({"a": int64s() | none(), "b": floats() | none()}))
-    def test_extra_field_but_none(self, *, data: StrMapping) -> None:
-        df = DataFrame(schema={"a": Int64, "b": Float64})
+    def test_disallow_missing_selected(self) -> None:
+        row = {"x": 3}
+        result = append_row(self.df, row, disallow_missing="x")
+        expected = DataFrame(
+            data=[*self.rows, (3, None)], schema=self.schema, orient="row"
+        )
+        assert_frame_equal(result, expected)
 
-        @dataclass(kw_only=True, slots=True)
-        class Row:
-            a: int | None = None
-            b: float | None = None
-            c: str | None = None
+    def test_disallow_null_selected(self) -> None:
+        row = {"x": 3, "y": None}
+        result = append_row(self.df, row, disallow_null="x")
+        expected = DataFrame(
+            data=[*self.rows, (3, None)], schema=self.schema, orient="row"
+        )
+        assert_frame_equal(result, expected)
 
-        row = Row(**data)
-        result = append_dataclass(df, row)
-        height = 0 if (row.a is None) and (row.b is None) else 1
-        check_polars_dataframe(result, height=height, schema_list=df.schema)
+    def test_in_place(self) -> None:
+        df = DataFrame(data=self.rows, schema=self.schema, orient="row")
+        assert df.height == 3
+        row = {"x": 3, "y": 6}
+        _ = append_row(df, row, in_place=True)
+        assert df.height == 4
 
-    @given(data=fixed_dictionaries({"datetime": zoned_date_times()}))
-    def test_zoned_datetime(self, *, data: StrMapping) -> None:
-        df = DataFrame(schema={"datetime": DatetimeUTC})
+    def test_error_predicate(self) -> None:
+        row = {"x": 3}
+        with raises(_AppendRowPredicateError, match="Predicate failed; got {'x': 3}"):
+            _ = append_row(self.df, row, predicate=lambda row: "y" in row)
 
-        @dataclass(kw_only=True, slots=True)
-        class Row:
-            datetime: ZonedDateTime
+    def test_error_disallow_extra(self) -> None:
+        row = {"x": 3, "y": 6, "z": None}
+        with raises(_AppendRowExtraKeysError, match=r"Extra key\(s\) found; got {'z'}"):
+            _ = append_row(self.df, row, disallow_extra=True)
 
-        row = Row(**data)
-        result = append_dataclass(df, row)
-        check_polars_dataframe(result, height=1, schema_list=df.schema)
+    def test_error_disallow_missing_all(self) -> None:
+        with raises(
+            _AppendRowMissingKeysError,
+            match=r"Missing key\(s\) found; got {'[xy]', '[xy]'}",
+        ):
+            _ = append_row(self.df, {}, disallow_missing=True)
 
-    @given(
-        data=fixed_dictionaries({
-            "a": int64s() | none(),
-            "b": floats() | none(),
-            "c": text_ascii(),
-        })
-    )
-    def test_error(self, *, data: StrMapping) -> None:
-        df = DataFrame(schema={"a": Int64, "b": Float64})
+    def test_error_disallow_missing_selected(self) -> None:
+        with raises(
+            _AppendRowMissingKeysError, match=r"Missing key\(s\) found; got {'x'}"
+        ):
+            _ = append_row(self.df, {}, disallow_missing="x")
 
-        @dataclass(kw_only=True, slots=True)
-        class Row:
-            a: int | None = None
-            b: float | None = None
-            c: str
+    def test_error_disallow_null_all(self) -> None:
+        row = {"x": None, "y": None}
+        with raises(
+            _AppendRowNullColumnsError,
+            match=r"Null column\(s\) found; got {'[xy]', '[xy]'}",
+        ):
+            _ = append_row(self.df, row, disallow_null=True)
 
-        row = Row(**data)
+    def test_error_disallow_null_selected(self) -> None:
+        row = {"x": None, "y": None}
         with raises(
-            AppendDataClassError,
-            match="Dataclass fields .* must be a subset of DataFrame columns .*; dataclass had extra items .*",
+            _AppendRowNullColumnsError, match=r"Null column\(s\) found; got {'x'}"
         ):
-            _ = append_dataclass(df, row)
+            _ = append_row(self.df, row, disallow_null="x")
 
 
 class TestAreFramesEqual:
@@ -2474,6 +2477,94 @@ class TestRoundToFloat:
             _ = round_to_float(cast("Any", x), cast("Any", y))
 
 
+class TestSearchPeriod:
+    @mark.parametrize(
+        ("time", "exp_start", "exp_end"),
+        [
+            param(whenever.Time(8, 58), None, None),
+            param(whenever.Time(8, 59), None, None),
+            param(whenever.Time(9), 0, None),
+            param(whenever.Time(9, 1), 0, 0),
+            param(whenever.Time(9, 59), 0, 0),
+            param(whenever.Time(10), 1, 0),
+            param(whenever.Time(10, 1), 1, 1),
+            param(whenever.Time(10, 59), 1, 1),
+            param(whenever.Time(11), 2, 1),
+            param(whenever.Time(11, 1), 2, 2),
+            param(whenever.Time(11, 59), 2, 2),
+            param(whenever.Time(12), None, 2),
+            param(whenever.Time(12, 1), None, None),
+            param(whenever.Time(12, 59), None, None),
+            param(whenever.Time(13), 3, None),
+            param(whenever.Time(13, 1), 3, 3),
+            param(whenever.Time(13, 59), 3, 3),
+            param(whenever.Time(14), None, 3),
+            param(whenever.Time(14, 1), None, None),
+            param(whenever.Time(14, 59), None, None),
+            param(whenever.Time(15), 4, None),
+            param(whenever.Time(15, 1), 4, 4),
+            param(whenever.Time(15, 59), 4, 4),
+            param(whenever.Time(16), None, 4),
+            param(whenever.Time(16, 1), None, None),
+            param(whenever.Time(16, 2), None, None),
+        ],
+    )
+    @mark.parametrize("start_or_end", [param("start"), param("end")])
+    def test_main(
+        self,
+        *,
+        time: whenever.Time,
+        start_or_end: Literal["start", "end"],
+        exp_start: int | None,
+        exp_end: int | None,
+    ) -> None:
+        date = whenever.Date(2000, 1, 1)
+        sr = DataFrame(
+            data=[
+                (
+                    date.at(whenever.Time(s)).py_datetime(),
+                    date.at(whenever.Time(e)).py_datetime(),
+                )
+                for s, e in [(9, 10), (10, 11), (11, 12), (13, 14), (15, 16)]
+            ],
+            schema={"start": DatetimeUTC, "end": DatetimeUTC},
+            orient="row",
+        ).with_columns(datetime=struct("start", "end"))["datetime"]
+        assert len(sr) == 5
+        date_time = date.at(time).assume_tz(UTC.key)
+        match start_or_end:
+            case "start":
+                expected = exp_start
+            case "end":
+                expected = exp_end
+            case never:
+                assert_never(never)
+        index = search_period(sr, date_time, start_or_end=start_or_end)
+        if expected is None:
+            assert index is None
+        else:
+            assert index is not None
+            assert 0 <= index <= (len(sr) - 1)
+            start, end = map(
+                to_zoned_date_time, cast("Iterable[dt.datetime]", sr[index].values())
+            )
+            match start_or_end:
+                case "start":
+                    assert start <= date_time < end
+                case "end":
+                    assert start < date_time <= end
+                case never:
+                    assert_never(never)
+            if index > 0:
+                prev_end = to_zoned_date_time(cast("dt.datetime", sr[index - 1]["end"]))
+                assert prev_end <= date_time
+            if index < (len(sr) - 1):
+                next_start = to_zoned_date_time(
+                    cast("dt.datetime", sr[index + 1]["start"])
+                )
+                assert date_time <= next_start
+
+
 class TestSelectExact:
     df: ClassVar[DataFrame] = DataFrame(
         data=[(True, False), (False, None), (None, True)],

{dycw_utilities-0.165.2 → dycw_utilities-0.166.0}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.165.2"
+__version__ = "0.166.0"

{dycw_utilities-0.165.2 → dycw_utilities-0.166.0}/src/utilities/polars.py

@@ -61,14 +61,12 @@ from utilities.gzip import read_binary
 from utilities.iterables import (
     CheckIterablesEqualError,
     CheckMappingsEqualError,
-    CheckSubSetError,
     CheckSuperMappingError,
     OneEmptyError,
     OneNonUniqueError,
     always_iterable,
     check_iterables_equal,
     check_mappings_equal,
-    check_subset,
     check_supermapping,
     is_iterable_not_str,
     one,
@@ -105,6 +103,7 @@ from utilities.whenever import (
 from utilities.zoneinfo import UTC, to_time_zone_name
 
 if TYPE_CHECKING:
+    import datetime as dt
     from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
     from collections.abc import Set as AbstractSet
 
@@ -300,29 +299,77 @@ def any_series(series: Series, /, *columns: ExprOrSeries) -> Series:
 ##
 
 
-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}
-    try:
-        check_subset(non_null_fields, df.columns)
-    except CheckSubSetError as error:
-        raise AppendDataClassError(
-            left=error.left, right=error.right, extra=error.extra
-        ) from None
-    row_cols = set(df.columns) & set(non_null_fields)
-    row = dataclass_to_dataframe(obj).select(*row_cols)
-    return concat([df, row], how="diagonal")
+def append_row(
+    df: DataFrame,
+    row: StrMapping,
+    /,
+    *,
+    predicate: Callable[[StrMapping], bool] | None = None,
+    disallow_extra: bool = False,
+    disallow_missing: bool | MaybeIterable[str] = False,
+    disallow_null: bool | MaybeIterable[str] = False,
+    in_place: bool = False,
+) -> DataFrame:
+    """Append a row to a DataFrame."""
+    if (predicate is not None) and not predicate(row):
+        raise _AppendRowPredicateError(df=df, row=row)
+    if disallow_extra and (len(extra := set(row) - set(df.columns)) >= 1):
+        raise _AppendRowExtraKeysError(df=df, row=row, extra=extra)
+    if disallow_missing is not False:
+        missing = set(df.columns) - set(row)
+        if disallow_missing is not True:
+            missing &= set(always_iterable(disallow_missing))
+        if len(missing) >= 1:
+            raise _AppendRowMissingKeysError(df=df, row=row, missing=missing)
+    other = DataFrame(data=[row], schema=df.schema)
+    if disallow_null:
+        other_null = other.select(col(c).is_null().any() for c in other.columns)
+        null = {k for k, v in other_null.row(0, named=True).items() if v}
+        if disallow_null is not True:
+            null &= set(always_iterable(disallow_null))
+        if len(null) >= 1:
+            raise _AppendRowNullColumnsError(df=df, row=row, columns=null)
+    return df.extend(other) if in_place else df.vstack(other)
+
+
+@dataclass(kw_only=True, slots=True)
+class AppendRowError(Exception):
+    df: DataFrame
+    row: StrMapping
+
+
+@dataclass(kw_only=True, slots=True)
+class _AppendRowPredicateError(AppendRowError):
+    @override
+    def __str__(self) -> str:
+        return f"Predicate failed; got {get_repr(self.row)}"
+
+
+@dataclass(kw_only=True, slots=True)
+class _AppendRowExtraKeysError(AppendRowError):
+    extra: AbstractSet[str]
+
+    @override
+    def __str__(self) -> str:
+        return f"Extra key(s) found; got {get_repr(self.extra)}"
+
+
+@dataclass(kw_only=True, slots=True)
+class _AppendRowMissingKeysError(AppendRowError):
+    missing: AbstractSet[str]
+
+    @override
+    def __str__(self) -> str:
+        return f"Missing key(s) found; got {get_repr(self.missing)}"
 
 
 @dataclass(kw_only=True, slots=True)
-class AppendDataClassError[T](Exception):
-    left: AbstractSet[T]
-    right: AbstractSet[T]
-    extra: AbstractSet[T]
+class _AppendRowNullColumnsError(AppendRowError):
+    columns: AbstractSet[str]
 
     @override
     def __str__(self) -> str:
-        return f"Dataclass fields {get_repr(self.left)} must be a subset of DataFrame columns {get_repr(self.right)}; dataclass had extra items {get_repr(self.extra)}"
+        return f"Null column(s) found; got {get_repr(self.columns)}"
 
 
 ##
@@ -2460,6 +2507,34 @@ class RoundToFloatError(Exception):
 ##
 
 
+def search_period(
+    series: Series,
+    date_time: ZonedDateTime,
+    /,
+    *,
+    start_or_end: Literal["start", "end"] = "end",
+) -> int | None:
+    """Search a series of periods for the one containing a given date-time."""
+    end = series.struct["end"]
+    py_date_time = date_time.py_datetime()
+    match start_or_end:
+        case "start":
+            index = end.search_sorted(py_date_time, side="right")
+            if index >= len(series):
+                return None
+            item: dt.datetime = series[index]["start"]
+            return index if py_date_time >= item else None
+        case "end":
+            index = end.search_sorted(py_date_time, side="left")
+            if index >= len(series):
+                return None
+            item: dt.datetime = series[index]["start"]
+            return index if py_date_time > item else None
+
+
+##
+
+
 def select_exact(
     df: DataFrame, /, *columns: IntoExprColumn, drop: MaybeIterable[str] | None = None
 ) -> DataFrame:
@@ -2672,6 +2747,7 @@ def zoned_date_time_period_dtype(
 
 
 __all__ = [
+    "AppendRowError",
     "BooleanValueCountsError",
     "CheckPolarsDataFrameError",
     "ColumnsToDictError",
@@ -2702,7 +2778,7 @@ __all__ = [
     "all_series",
     "any_dataframe_columns",
     "any_series",
-    "append_dataclass",
+    "append_row",
     "are_frames_equal",
     "bernoulli",
     "boolean_value_counts",
@@ -2750,6 +2826,7 @@ __all__ = [
     "read_series",
     "replace_time_zone",
     "round_to_float",
+    "search_period",
     "select_exact",
     "serialize_dataframe",
     "set_first_row_as_columns",