dycw-utilities 0.153.2__tar.gz → 0.153.3__tar.gz

{dycw_utilities-0.153.2 → dycw_utilities-0.153.3}/PKG-INFO

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

{dycw_utilities-0.153.2 → dycw_utilities-0.153.3}/pyproject.toml

@@ -190,7 +190,7 @@ slack-sdk-test = [
   "aiohttp",
 ]
 sqlalchemy = [
-  "sqlalchemy >=2.0.41, <2.1",
+  "sqlalchemy >=2.0.42, <2.1",
   "psycopg",
 ]
 sqlalchemy-polars = [
@@ -229,7 +229,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.153.2"
+version = "0.153.3"
 
 [project.entry-points.pytest11]
 pytest-randomly = "utilities.pytest_plugins.pytest_randomly"
@@ -262,7 +262,7 @@ test = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.153.2"
+current_version = "0.153.3"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.153.2 → dycw_utilities-0.153.3}/src/tests/test_polars.py

@@ -58,7 +58,7 @@ from polars.exceptions import ComputeError
 from polars.schema import Schema
 from polars.testing import assert_frame_equal, assert_series_equal
 from pytest import mark, param, raises
-from whenever import TimeZoneNotFoundError
+from whenever import Time, TimeZoneNotFoundError, ZonedDateTime
 
 import utilities.polars
 from utilities.hypothesis import (
@@ -173,7 +173,9 @@ from utilities.polars import (
     nan_sum_agg,
     nan_sum_cols,
     normal,
+    offset_datetime,
     order_of_magnitude,
+    period_range,
     read_dataframe,
     read_series,
     reify_exprs,
@@ -2205,6 +2207,15 @@ class TestNormal:
         assert series.is_finite().all()
 
 
+class TestOffsetDateTime:
+    @mark.parametrize(("n", "time"), [param(1, Time(13, 30)), param(2, Time(15))])
+    def test_main(self, *, n: int, time: Time) -> None:
+        datetime = ZonedDateTime(2000, 1, 1, 12, tz=UTC.key)
+        result = offset_datetime(datetime, "1h30m", n=n)
+        expected = datetime.replace_time(time)
+        assert result == expected
+
+
 class TestOrderOfMagnitude:
     @given(
         sign=sampled_from([1, -1]),
@@ -2231,6 +2242,19 @@ class TestOrderOfMagnitude:
         assert (res_int == exp_int).all()
 
 
+class TestPeriodRange:
+    start: ClassVar[ZonedDateTime] = ZonedDateTime(2000, 1, 1, 12, tz=UTC.key)
+    end: ClassVar[ZonedDateTime] = ZonedDateTime(2000, 1, 1, 15, tz=UTC.key)
+
+    @mark.parametrize("end_or_length", [param(end), param(3)])
+    def test_main(self, *, end_or_length: ZonedDateTime | int) -> None:
+        rng = period_range(self.start, end_or_length, interval="1h", eager=True)
+        assert len(rng) == 3
+        assert rng.dtype == zoned_datetime_period_dtype()
+        assert rng[0]["start"] == self.start.py_datetime()
+        assert rng[-1]["end"] == self.end.py_datetime()
+
+
 class TestReifyExprs:
     @given(length=hypothesis.strategies.integers(0, 10), name=text_ascii())
     def test_one_expr(self, *, length: int, name: str) -> None:

{dycw_utilities-0.153.2 → dycw_utilities-0.153.3}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.153.2"
+__version__ = "0.153.3"

{dycw_utilities-0.153.2 → dycw_utilities-0.153.3}/src/utilities/polars.py

@@ -33,6 +33,7 @@ from polars import (
     any_horizontal,
     col,
     concat,
+    datetime_range,
     int_range,
     lit,
     struct,
@@ -49,6 +50,7 @@ from polars.exceptions import (
 )
 from polars.schema import Schema
 from polars.testing import assert_frame_equal, assert_series_equal
+from whenever import ZonedDateTime
 
 from utilities.dataclasses import _YieldFieldsInstance, yield_fields
 from utilities.errors import ImpossibleCaseError
@@ -129,6 +131,7 @@ DatetimeUSEastern = Datetime(time_zone="US/Eastern")
 DatetimeUTC = Datetime(time_zone="UTC")
 _FINITE_EWM_MIN_WEIGHT = 0.9999
 
+
 ##
 
 
@@ -2021,6 +2024,19 @@ def normal(
 ##
 
 
+def offset_datetime(
+    datetime: ZonedDateTime, offset: str, /, *, n: int = 1
+) -> ZonedDateTime:
+    """Offset a datetime as `polars` would."""
+    sr = Series(values=[datetime.py_datetime()])
+    for _ in range(n):
+        sr = sr.dt.offset_by(offset)
+    return ZonedDateTime.from_py_datetime(sr.item())
+
+
+##
+
+
 @overload
 def order_of_magnitude(column: ExprLike, /, *, round_: bool = False) -> Expr: ...
 @overload
@@ -2041,6 +2057,75 @@ def order_of_magnitude(
 ##
 
 
+@overload
+def period_range(
+    start: ZonedDateTime,
+    end_or_length: ZonedDateTime | int,
+    /,
+    *,
+    interval: str = "1d",
+    time_unit: TimeUnit | None = None,
+    time_zone: TimeZoneLike | None = None,
+    eager: Literal[True],
+) -> Series: ...
+@overload
+def period_range(
+    start: ZonedDateTime,
+    end_or_length: ZonedDateTime | int,
+    /,
+    *,
+    interval: str = "1d",
+    time_unit: TimeUnit | None = None,
+    time_zone: TimeZoneLike | None = None,
+    eager: Literal[False] = False,
+) -> Expr: ...
+@overload
+def period_range(
+    start: ZonedDateTime,
+    end_or_length: ZonedDateTime | int,
+    /,
+    *,
+    interval: str = "1d",
+    time_unit: TimeUnit | None = None,
+    time_zone: TimeZoneLike | None = None,
+    eager: bool = False,
+) -> Series | Expr: ...
+def period_range(
+    start: ZonedDateTime,
+    end_or_length: ZonedDateTime | int,
+    /,
+    *,
+    interval: str = "1d",
+    time_unit: TimeUnit | None = None,
+    time_zone: TimeZoneLike | None = None,
+    eager: bool = False,
+) -> Series | Expr:
+    """Construct a period range."""
+    time_zone_use = None if time_zone is None else ensure_time_zone(time_zone).key
+    match end_or_length:
+        case ZonedDateTime() as end:
+            ...
+        case int() as length:
+            end = offset_datetime(start, interval, n=length)
+        case never:
+            assert_never(never)
+    starts = datetime_range(
+        start.py_datetime(),
+        end.py_datetime(),
+        interval,
+        closed="left",
+        time_unit=time_unit,
+        time_zone=time_zone_use,
+        eager=eager,
+    ).alias("start")
+    ends = (starts.dt.offset_by(interval)).alias("end")
+    period = struct(starts, ends)
+    return try_reify_expr(period, starts, ends)
+
+
+##
+
+
 def reify_exprs(
     *exprs: IntoExprColumn, **named_exprs: IntoExprColumn
 ) -> Expr | Series | DataFrame:
@@ -2529,7 +2614,9 @@ __all__ = [
     "nan_sum_agg",
     "nan_sum_cols",
     "normal",
+    "offset_datetime",
     "order_of_magnitude",
+    "period_range",
     "read_dataframe",
     "read_series",
     "replace_time_zone",