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

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

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

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

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

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

@@ -194,6 +194,7 @@ from utilities.polars import (
     reify_exprs,
     replace_time_zone,
     round_to_float,
+    search_period,
     select_exact,
     serialize_dataframe,
     serialize_series,
@@ -224,6 +225,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
 
@@ -2474,6 +2476,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.165.3}/src/utilities/__init__.py

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

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

@@ -105,6 +105,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
 
@@ -2460,6 +2461,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."""
+    start, end = [series.struct[k] for k in ["start", "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:
@@ -2750,6 +2779,7 @@ __all__ = [
     "read_series",
     "replace_time_zone",
     "round_to_float",
+    "search_period",
     "select_exact",
     "serialize_dataframe",
     "set_first_row_as_columns",