Add dataclass_to_row

src/tests/test_polars.py

@@ -7,7 +7,14 @@
 from zoneinfo import ZoneInfo
 
 from hypothesis import given
-from hypothesis.strategies import fixed_dictionaries, floats, none, sampled_from
+from hypothesis.strategies import (
+    dates,
+    datetimes,
+    fixed_dictionaries,
+    floats,
+    none,
+    sampled_from,
+)
 from polars import (
     Boolean,
     DataFrame,
@@ -29,7 +36,7 @@
 from polars.testing import assert_frame_equal, assert_series_equal
 from pytest import mark, param, raises
 
-from utilities.hypothesis import int64s, text_ascii
+from utilities.hypothesis import int64s, text_ascii, zoned_datetimes
 from utilities.polars import (
     AppendDataClassError,
     CheckPolarsDataFrameError,
@@ -58,6 +65,7 @@
     collect_series,
     columns_to_dict,
     convert_time_zone,
+    dataclass_to_row,
     drop_null_struct_series,
     ensure_expr_or_series,
     floor_datetime,
@@ -99,6 +107,19 @@ class Row:
         result = append_dataclass(df, row)
         check_polars_dataframe(result, height=1, schema_list=schema)
 
+    @given(data=fixed_dictionaries({"datetime": zoned_datetimes()}))
+    def test_zoned_datetime(self, *, data: StrMapping) -> None:
+        schema = {"datetime": DatetimeUTC}
+        df = DataFrame([], schema=schema, orient="row")
+
+        @dataclass(kw_only=True)
+        class Row:
+            datetime: dt.datetime
+
+        row = Row(**data)
+        result = append_dataclass(df, row)
+        check_polars_dataframe(result, height=1, schema_list=schema)
+
     @given(
         data=fixed_dictionaries({
             "a": int64s() | none(),
@@ -548,6 +569,78 @@ def test_dataframe_nested_twice(self) -> None:
         assert_frame_equal(result, expected)
 
 
+class TestDataClassToRow:
+    @given(
+        data=fixed_dictionaries({
+            "a": int64s() | none(),
+            "b": floats() | none(),
+            "c": dates() | none(),
+            "d": datetimes() | none(),
+        })
+    )
+    def test_basic_types(self, *, data: StrMapping) -> None:
+        @dataclass(kw_only=True)
+        class Row:
+            a: int | None = None
+            b: float | None = None
+            c: dt.date | None = None
+            d: dt.datetime | None = None
+
+        df = dataclass_to_row(Row(**data))
+        check_polars_dataframe(
+            df,
+            height=1,
+            schema_list={"a": Int64, "b": Float64, "c": Date, "d": Datetime},
+        )
+
+    @given(data=fixed_dictionaries({"datetime": zoned_datetimes()}))
+    def test_zoned_datetime(self, *, data: StrMapping) -> None:
+        @dataclass(kw_only=True)
+        class Row:
+            datetime: dt.datetime
+
+        row = Row(**data)
+        df = dataclass_to_row(row)
+        check_polars_dataframe(df, height=1, schema_list={"datetime": DatetimeUTC})
+
+    @given(
+        data=fixed_dictionaries({
+            "a": int64s(),
+            "b": int64s(),
+            "inner": fixed_dictionaries({
+                "start": zoned_datetimes(),
+                "end": zoned_datetimes(),
+            }),
+        })
+    )
+    def test_zoned_datetime_nested(self, *, data: StrMapping) -> None:
+        @dataclass(kw_only=True)
+        class Inner:
+            start: dt.datetime
+            end: dt.datetime
+
+        inner = Inner(**data["inner"])
+
+        @dataclass(kw_only=True)
+        class Outer:
+            a: int | None = None
+            b: int | None = None
+            inner: Inner | None = None
+
+        data = dict(data) | {"inner": inner}
+        outer = Outer(**data)
+        df = dataclass_to_row(outer)
+        check_polars_dataframe(
+            df,
+            height=1,
+            schema_list={
+                "a": Int64,
+                "b": Int64,
+                "inner": Struct({"start": DatetimeUTC, "end": DatetimeUTC}),
+            },
+        )
+
+
 class TestDatetimeUTC:
     @mark.parametrize(
         ("dtype", "time_zone"),

src/tests/test_redis.py

@@ -26,6 +26,21 @@ async def test_main(self, *, data: DataObject, value: bool) -> None:
                     _ = await key.set_async(value, db=15)
                     assert await key.get_async(db=15) is value
 
+    @given(data=data(), value=booleans())
+    @SKIPIF_CI_AND_NOT_LINUX
+    async def test_using_client(self, *, data: DataObject, value: bool) -> None:
+        async with redis_cms(data) as container:
+            key = RedisKey(name=container.key, type=bool)
+            match container.client:
+                case redis.Redis() as client:
+                    assert key.get(client=client) is None
+                    _ = key.set(value, client=client)
+                    assert key.get(client=client) is value
+                case redis.asyncio.Redis() as client:
+                    assert await key.get_async(client=client) is None
+                    _ = await key.set_async(value, client=client)
+                    assert await key.get_async(client=client) is value
+
 
 class TestRedisHashMapKey:
     @given(data=data(), key=int64s(), value=booleans())
@@ -42,3 +57,20 @@ async def test_main(self, *, data: DataObject, key: int, value: bool) -> None:
                     assert await hash_map_key.hget_async(key, db=15) is None
                     _ = await hash_map_key.hset_async(key, value, db=15)
                     assert await hash_map_key.hget_async(key, db=15) is value
+
+    @given(data=data(), key=int64s(), value=booleans())
+    @SKIPIF_CI_AND_NOT_LINUX
+    async def test_using_client(
+        self, *, data: DataObject, key: int, value: bool
+    ) -> None:
+        async with redis_cms(data) as container:
+            hash_map_key = RedisHashMapKey(name=container.key, key=int, value=bool)
+            match container.client:
+                case redis.Redis() as client:
+                    assert hash_map_key.hget(key, client=client) is None
+                    _ = hash_map_key.hset(key, value, client=client)
+                    assert hash_map_key.hget(key, client=client) is value
+                case redis.asyncio.Redis() as client:
+                    assert await hash_map_key.hget_async(key, client=client) is None
+                    _ = await hash_map_key.hset_async(key, value, client=client)
+                    assert await hash_map_key.hget_async(key, client=client) is value

src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.57.2"
+__version__ = "0.57.3"

src/utilities/polars.py

@@ -8,7 +8,7 @@
 from dataclasses import dataclass
 from datetime import timezone
 from enum import Enum
-from functools import reduce
+from functools import partial, reduce
 from itertools import chain
 from typing import (
     TYPE_CHECKING,
@@ -73,6 +73,7 @@
     one,
 )
 from utilities.math import CheckIntegerError, check_integer
+from utilities.types import ensure_datetime
 from utilities.typing import (
     get_args,
     is_frozenset_type,
@@ -81,7 +82,7 @@
     is_optional_type,
     is_set_type,
 )
-from utilities.zoneinfo import UTC, get_time_zone_name
+from utilities.zoneinfo import UTC, ensure_time_zone, get_time_zone_name
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Iterator, Sequence
@@ -106,7 +107,7 @@ def append_dataclass(df: DataFrame, obj: Dataclass, /) -> DataFrame:
         raise AppendDataClassError(
             left=error.left, right=error.right, extra=error.extra
         ) from None
-    row = DataFrame([obj], orient="row")
+    row = dataclass_to_row(obj)
     return concat([df, row], how="diagonal")
 
 
@@ -506,6 +507,28 @@ def __str__(self) -> str:
         return f"DataFrame must be unique on {self.key!r}\n\n{self.df}"
 
 
+def dataclass_to_row(obj: Dataclass, /) -> DataFrame:
+    """Convert a dataclass into a 1-row DataFrame."""
+    df = DataFrame([obj], orient="row")
+    return reduce(partial(_dataclass_to_row_reducer, obj=obj), df.columns, df)
+
+
+def _dataclass_to_row_reducer(
+    df: DataFrame, column: str, /, *, obj: Dataclass
+) -> DataFrame:
+    dtype = df[column].dtype
+    if isinstance(dtype, Datetime):
+        datetime = ensure_datetime(getattr(obj, column), nullable=True)
+        if (datetime is None) or (datetime.tzinfo is None):
+            return df
+        time_zone = ensure_time_zone(datetime.tzinfo)
+        return df.with_columns(col(column).cast(zoned_datetime(time_zone=time_zone)))
+    if isinstance(dtype, Struct):
+        inner = dataclass_to_row(getattr(obj, column))
+        return df.with_columns(**{column: inner.select(all=struct("*"))["all"]})
+    return df
+
+
 def drop_null_struct_series(series: Series, /) -> Series:
     """Drop nulls in a struct-dtype Series as per the <= 1.1 definition."""
     try:
@@ -900,6 +923,7 @@ def zoned_datetime(
     "collect_series",
     "columns_to_dict",
     "convert_time_zone",
+    "dataclass_to_row",
     "drop_null_struct_series",
     "ensure_expr_or_series",
     "floor_datetime",