Repository: unionai-oss/pandera

Issue:
Time-agnostic DateTime with pandera-native polars datatype using DataFrameModel not working
The use of `dtype_kwargs` for the `pandera.engines.polars_engine.DateTime` dtype as demonstrated in the [documentation examples](https://pandera.readthedocs.io/en/stable/polars.html#time-agnostic-datetime) does not seem to work when using with a `DataFrameModel`.

- [x] I have checked that this issue has not already been reported.
- [x] I have confirmed this bug exists on the latest version of pandera -> using 0.19.2.
- [ ] (optional) I have confirmed this bug exists on the main branch of pandera.


I tried with all 3 variations of the `DataFrameModel`, as well as using a `DataFrameSchema` and only the latter seems to work as expected:

#### Code Sample:

```python

import datetime as dt

import pandera.polars as pa
import polars as pl
from pandera.engines import polars_engine as pe

df = pl.DataFrame(
    schema={
        "column_1": pl.Utf8,
        "column_2": pl.Datetime(time_zone="UTC"),
    },
    data={
        "column_1": ["value1", "value2"],
        "column_2": [dt.datetime(2024, 1, 1), dt.datetime(2024, 1, 2)],
    },
)


class SchemaDataFrameModel1(pa.DataFrameModel):
    column_1: str
    column_2: dt.datetime


class SchemaDataFrameModel2(pa.DataFrameModel):
    column_1: str
    column_2: pe.DateTime = pa.Field(dtype_kwargs={"time_zone_agnostic": True})


class SchemaDataFrameModel3(pa.DataFrameModel):
    column_1: str
    column_2: Annotated[pe.DateTime, True]


schema_dataframeschema_1 = pa.DataFrameSchema(
    {
        "column_1": pa.Column(str),
        "column_2": pa.Column(dt.datetime),
    }
)


schema_dataframeschema_2 = pa.DataFrameSchema(
    {
        "column_1": pa.Column(str),
        "column_2": pa.Column(pe.DateTime(time_zone_agnostic=True)),
    }
)


cases = {
    "DataFrameModel (Field) - without `time_zone_agnostic=True`": SchemaDataFrameModel1,
    "DataFrameModel (Field) - with `time_zone_agnostic=True": SchemaDataFrameModel2,
    "DataFrameModel (Annotated)": SchemaDataFrameModel3,
    "DataFrameSchema - without `time_zone_agnostic=True`": schema_dataframeschema_1,
    "DataFrameSchema - with `time_zone_agnostic=True`": schema_dataframeschema_2,
}

for case, schema in cases.items():
    print(f"Case: {case}")
    try:
        if type(schema) == pa.DataFrameModel:
            schema.to_schema().validate(df)
        else:
            schema.validate(df)
        print("\t✅ Validation successful")
    except Exception as e:
        print(f"\t❌ Validation Failed: {e}")


```

#### Output:
```
Case: DataFrameModel (Field) - without `time_zone_agnostic=True`
        ❌ Validation Failed: expected column 'column_2' to have type Datetime(time_unit='us', time_zone=None), got Datetime(time_unit='us', time_zone='UTC')
Case: DataFrameModel (Field) - with `time_zone_agnostic=True
        ❌ Validation Failed: 'Datetime' object is not callable
Case: DataFrameModel (Annotated)
        ❌ Validation Failed: Annotation 'DateTime' requires all positional arguments ['time_zone_agnostic', 'time_zone', 'time_unit'].
Case: DataFrameSchema - without `time_zone_agnostic=True`
        ❌ Validation Failed: expected column 'column_2' to have type Datetime(time_unit='us', time_zone=None), got Datetime(time_unit='us', time_zone='UTC')
Case: DataFrameSchema - with `time_zone_agnostic=True`
        ✅ Validation successful
```

#### Expected behaviour
I expect the validation to fail on schemas that don't provide `time_zone_agnostic=True` (which is the case), and for it to pass validation when setting `time_zone_agnostic=True`.

#### Actual behaviour
The use of `pa.Field` fails with `'Datetime' object is not callable` and the use of `Annotated` fails with `Annotation 'DateTime' requires all positional arguments ['time_zone_agnostic', 'time_zone', 'time_unit']`.

For the case of `pa.Field`, it looks like an instance of `pl.Datetime` gets returned by `engine_dtype = pe.Engine.dtype(annotation.raw_annotation)` in `_build_columns()` of `class DataFrameModel`, and then called again with `dtype(**self.dtype_kwargs)` in `_get_schema_properties()` of `class FieldInfo` which throws the error.

#### Desktop

 - OS: Ubuntu 22.04 LTS (WSL)


Discussion:

Resolving patch:
diff --git a/pandera/api/polars/model.py b/pandera/api/polars/model.py
--- a/pandera/api/polars/model.py
+++ b/pandera/api/polars/model.py
@@ -1,5 +1,6 @@
 """Class-based api for polars models."""
 
+import inspect
 from typing import Dict, List, Tuple, Type
 
 import pandas as pd
@@ -47,10 +48,16 @@ def _build_columns(  # pylint:disable=too-many-locals
             field_name = field.name
             check_name = getattr(field, "check_name", None)
 
-            engine_dtype = None
             try:
                 engine_dtype = pe.Engine.dtype(annotation.raw_annotation)
-                dtype = engine_dtype.type
+                if inspect.isclass(annotation.raw_annotation) and issubclass(
+                    annotation.raw_annotation, pe.DataType
+                ):
+                    # use the raw annotation as the dtype if it's a native
+                    # pandera polars datatype
+                    dtype = annotation.raw_annotation
+                else:
+                    dtype = engine_dtype.type
             except (TypeError, ValueError) as exc:
                 if annotation.metadata:
                     if field.dtype_kwargs:
@@ -64,13 +71,13 @@ def _build_columns(  # pylint:disable=too-many-locals
                 elif annotation.default_dtype:
                     dtype = annotation.default_dtype
                 else:
-                    dtype = annotation.arg
+                    dtype = annotation.arg  # type: ignore
 
             if (
                 annotation.origin is None
                 or isinstance(annotation.origin, pl.datatypes.DataTypeClass)
                 or annotation.origin is Series
-                or engine_dtype
+                or dtype
             ):
                 if check_name is False:
                     raise SchemaInitError(


Test patch:
diff --git a/tests/polars/test_polars_model.py b/tests/polars/test_polars_model.py
--- a/tests/polars/test_polars_model.py
+++ b/tests/polars/test_polars_model.py
@@ -3,9 +3,18 @@
 import sys
 from typing import Optional
 
+try:  # python 3.9+
+    from typing import Annotated  # type: ignore
+except ImportError:
+    from typing_extensions import Annotated  # type: ignore
+
 import polars as pl
 import pytest
+from hypothesis import given
+from hypothesis import strategies as st
+from polars.testing.parametric import column, dataframes
 
+import pandera.engines.polars_engine as pe
 from pandera.errors import SchemaError
 from pandera.polars import (
     Column,
@@ -211,3 +220,53 @@ class ModelWithNestedDtypes(DataFrameModel):
 
     schema = ModelWithNestedDtypes.to_schema()
     assert schema_with_list_type == schema
+
+
+@pytest.mark.parametrize(
+    "time_zone",
+    [
+        None,
+        "UTC",
+        "GMT",
+        "EST",
+    ],
+)
+@given(st.data())
+def test_dataframe_schema_with_tz_agnostic_dates(time_zone, data):
+    strategy = dataframes(
+        column("datetime_col", dtype=pl.Datetime()),
+        lazy=True,
+        size=10,
+    )
+    lf = data.draw(strategy)
+    lf = lf.cast({"datetime_col": pl.Datetime(time_zone=time_zone)})
+
+    class ModelTZAgnosticKwargs(DataFrameModel):
+        datetime_col: pe.DateTime = Field(
+            dtype_kwargs={"time_zone_agnostic": True}
+        )
+
+    class ModelTZSensitiveKwargs(DataFrameModel):
+        datetime_col: pe.DateTime = Field(
+            dtype_kwargs={"time_zone_agnostic": False}
+        )
+
+    class ModelTZAgnosticAnnotated(DataFrameModel):
+        datetime_col: Annotated[pe.DateTime, True, "us", None]
+
+    class ModelTZSensitiveAnnotated(DataFrameModel):
+        datetime_col: Annotated[pe.DateTime, False, "us", None]
+
+    for tz_agnostic_model in (
+        ModelTZAgnosticKwargs,
+        ModelTZAgnosticAnnotated,
+    ):
+        tz_agnostic_model.validate(lf)
+
+    for tz_sensitive_model in (
+        ModelTZSensitiveKwargs,
+        ModelTZSensitiveAnnotated,
+    ):
+        if time_zone:
+            with pytest.raises(SchemaError):
+                tz_sensitive_model.validate(lf)