switch to using dtype validation for datetime timezones

Author: dfrankland

dataframely/columns/datetime.py

@@ -273,7 +273,7 @@ def __init__(
         max: dt.datetime | None = None,
         max_exclusive: dt.datetime | None = None,
         resolution: str | None = None,
-        time_zone: ZoneInfo | str | dt.timezone | None = None,
+        time_zone: str | dt.tzinfo | None = None,
         check: Callable[[pl.Expr], pl.Expr] | None = None,
         alias: str | None = None,
         metadata: dict[str, Any] | None = None,
@@ -339,16 +339,6 @@ def validation_rules(self, expr: pl.Expr) -> dict[str, pl.Expr]:
         result = super().validation_rules(expr)
         if self.resolution is not None:
             result["resolution"] = expr.dt.truncate(self.resolution) == expr
-        if self.time_zone is not None:
-            time_zone = (
-                self.time_zone.key
-                if isinstance(self.time_zone, ZoneInfo)
-                else self.time_zone
-            )
-            result["time_zone"] = pl.coalesce(
-                expr == pl.selectors.datetime(time_unit="us", time_zone=time_zone),
-                False,
-            )
         return result
 
     def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:

dataframely/random.py

@@ -4,7 +4,6 @@
 import datetime as dt
 from collections.abc import Sequence
 from typing import TypeVar
-from zoneinfo import ZoneInfo
 
 import numpy as np
 import polars as pl
@@ -294,7 +293,7 @@ def sample_datetime(
         min: dt.datetime,
         max: dt.datetime | None,
         resolution: str | None = None,
-        time_zone: ZoneInfo | str | dt.timezone | None = None,
+        time_zone: str | dt.tzinfo | None = None,
         null_probability: float = 0.0,
     ) -> pl.Series:
         """Sample a list of datetimes in the provided range.

tests/column_types/test_datetime.py

@@ -2,15 +2,16 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import datetime as dt
+import re
 from typing import Any
-from zoneinfo import ZoneInfo
 
 import polars as pl
 import pytest
 from polars.testing import assert_frame_equal
 
 import dataframely as dy
 from dataframely.columns import Column
+from dataframely.exc import DtypeValidationError
 from dataframely.random import Generator
 from dataframely.testing import evaluate_rules, rules_from_exprs
 from dataframely.testing.factory import create_schema
@@ -408,122 +409,46 @@ def test_validate_resolution(
 
 
 @pytest.mark.parametrize(
-    ("column", "values", "schema", "valid"),
+    "column",
     [
-        (
-            dy.Datetime(),
-            [
-                dt.datetime(2020, 4, 5, tzinfo=dt.UTC),
-                dt.datetime(2021, 1, 1, 12, tzinfo=dt.UTC),
-                dt.datetime(2022, 7, 10, 0, 0, 1, tzinfo=dt.UTC),
-            ],
-            None,
-            {},
-        ),
-        (
-            dy.Datetime(time_zone="UTC"),
-            [
-                dt.datetime(2020, 4, 5),
-                dt.datetime(2021, 1, 1, 12),
-                dt.datetime(2022, 7, 10, 0, 0, 1),
-            ],
-            None,
-            {"time_zone": [False, False, False]},
-        ),
-        (
-            dy.Datetime(time_zone="UTC"),
-            [
-                dt.datetime(2020, 4, 5, tzinfo=dt.UTC),
-                dt.datetime(2021, 1, 1, 12, tzinfo=dt.UTC),
-                dt.datetime(2022, 7, 10, 0, 0, 1, tzinfo=dt.UTC),
-            ],
-            None,
-            {"time_zone": [True, True, True]},
-        ),
-        (
-            dy.Datetime(
-                time_zone=dt.timezone(
-                    dt.timedelta(hours=-7), name="America/Los_Angeles"
-                )
-            ),
-            [
-                dt.datetime(2020, 4, 5),
-                dt.datetime(2021, 1, 1, 12),
-                dt.datetime(2022, 7, 10, 0, 0, 1),
-            ],
-            pl.Datetime(time_zone="America/Los_Angeles"),
-            {"time_zone": [True, True, True]},
-        ),
-        (
-            dy.Datetime(
-                time_zone=dt.timezone(
-                    dt.timedelta(hours=-7), name="America/Los_Angeles"
-                )
-            ),
-            [
-                dt.datetime(2020, 4, 5),
-                dt.datetime(2021, 1, 1, 12),
-                dt.datetime(2022, 7, 10, 0, 0, 1),
-            ],
-            None,
-            {"time_zone": [False, False, False]},
-        ),
-        (
-            dy.Datetime(time_zone=dt.timezone(dt.timedelta(hours=-7))),
-            [
-                dt.datetime(2020, 4, 5),
-                dt.datetime(2021, 1, 1, 12),
-                dt.datetime(2022, 7, 10, 0, 0, 1),
-            ],
-            pl.Datetime(time_zone="America/Los_Angeles"),
-            {"time_zone": [False, False, False]},
-        ),
-        (
-            dy.Datetime(time_zone=ZoneInfo("Etc/UTC")),
-            [
-                dt.datetime(2020, 4, 5, tzinfo=ZoneInfo("Etc/UTC")),
-                dt.datetime(2021, 1, 1, 12, tzinfo=ZoneInfo("Etc/UTC")),
-                dt.datetime(2022, 7, 10, 0, 0, 1, tzinfo=ZoneInfo("Etc/UTC")),
-            ],
-            None,
-            {"time_zone": [True, True, True]},
-        ),
-        (
-            dy.Datetime(time_zone=ZoneInfo("Etc/UTC")),
-            [
-                dt.datetime(2020, 4, 5, tzinfo=ZoneInfo("America/New_York")),
-                dt.datetime(2021, 1, 1, 12, tzinfo=ZoneInfo("America/New_York")),
-                dt.datetime(2022, 7, 10, 0, 0, 1, tzinfo=ZoneInfo("America/New_York")),
-            ],
-            None,
-            {"time_zone": [False, False, False]},
+        dy.Datetime(
+            min=dt.datetime(2020, 1, 1), max=dt.datetime(2021, 1, 1), resolution="1h"
         ),
+        dy.Datetime(time_zone="Etc/UTC"),
     ],
 )
-def test_validate_datetime_timezone(
-    column: Column,
-    values: list[Any],
-    schema: pl.Datetime | None,
-    valid: dict[str, list[bool]],
-) -> None:
-    lf = pl.LazyFrame(
-        {"a": values}, schema={"a": schema} if schema is not None else None
-    )
-    actual = evaluate_rules(lf, rules_from_exprs(column.validation_rules(pl.col("a"))))
-    expected = pl.LazyFrame(valid)
-    assert_frame_equal(actual, expected)
+def test_sample(column: dy.Column) -> None:
+    generator = Generator(seed=42)
+    samples = column.sample(generator, n=10_000)
+    schema = create_schema("test", {"a": column})
+    schema.validate(samples.to_frame("a"))
 
 
 @pytest.mark.parametrize(
-    "column",
+    ("dtype", "column", "error"),
     [
-        dy.Datetime(
-            min=dt.datetime(2020, 1, 1), max=dt.datetime(2021, 1, 1), resolution="1h"
-        )
+        (
+            pl.Datetime(time_zone="America/New_York"),
+            dy.Datetime(time_zone="Etc/UTC"),
+            r"1 columns have an invalid dtype.*\n.*got dtype 'Datetime\(time_unit='us', time_zone='America/New_York'\)' but expected 'Datetime\(time_unit='us', time_zone='Etc/UTC'\)'",
+        ),
+        (
+            pl.Datetime(time_zone="Etc/UTC"),
+            dy.Datetime(time_zone="Etc/UTC"),
+            None,
+        ),
     ],
 )
-def test_sample_resolution(column: dy.Column) -> None:
-    generator = Generator(seed=42)
-    samples = column.sample(generator, n=10_000)
+def test_dtype_validation(
+    dtype: pl.DataType,
+    column: dy.Column,
+    error: str | None,
+) -> None:
+    df = pl.DataFrame(schema={"a": dtype})
     schema = create_schema("test", {"a": column})
-    schema.validate(samples.to_frame("a"))
+    if error is None:
+        schema.validate(df)
+    else:
+        with pytest.raises(DtypeValidationError) as exc:
+            schema.validate(df)
+        assert re.match(error, str(exc.value))