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,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import datetime as dt
+import re
 from typing import Any
 from zoneinfo import ZoneInfo
 
@@ -11,6 +12,7 @@
 
 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
@@ -519,11 +521,42 @@ def test_validate_datetime_timezone(
     [
         dy.Datetime(
             min=dt.datetime(2020, 1, 1), max=dt.datetime(2021, 1, 1), resolution="1h"
-        )
+        ),
+        dy.Datetime(time_zone="Etc/UTC"),
     ],
 )
-def test_sample_resolution(column: dy.Column) -> None:
+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(
+    ("dtype", "column", "error"),
+    [
+        (
+            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_dtype_validation(
+    dtype: pl.DataType,
+    column: dy.Column,
+    error: str | None,
+) -> None:
+    df = pl.DataFrame(schema={"a": dtype})
+    schema = create_schema("test", {"a": column})
+    if error is None:
+        schema.validate(df)
+    else:
+        with pytest.raises(DtypeValidationError) as exc:
+            schema.validate(df)
+        assert re.match(error, str(exc.value))