Repository: unionai-oss/pandera

Issue:
Pandas deprecated `(timezone naive timestamp).astype(timezone aware dtype)`
Pandas support for `(timezone naive timestamp).astype(timezone aware dtype)` is deprecated and will be removed in future versions.
This line https://github.com/unionai-oss/pandera/blob/master/pandera/engines/pandas_engine.py#L740 gives a warning when coercing a timezone naive series to a timezone aware one with a schema like the one below:
```python
class MySchema(SchemaModel):
    col : Series[DatetimeTZDType] = Field(coerce=True, dtype_kwargs={'tz': 'utc'}
```

Example code:
```python
from pandas import DataFrame, DatetimeTZDtype, to_datetime

from pandera import Field, SchemaModel
from pandera.dtypes import DateTime
from pandera.typing import Series


class TimeZoneSchema(SchemaModel):
    column: Series[DatetimeTZDtype] = Field(coerce=True, dtype_kwargs={'tz': 'UTC'})


df = DataFrame({'column': to_datetime(['2022-04-30', '2022-05-01', '2022-05-02'])})
TimeZoneSchema.validate(df)
```

example output:

```text
D:\Miniconda3\envs\pandera-types\lib\site-packages\pandera\engines\pandas_engine.py:740: FutureWarning: Using .astype to convert from timezone-naive dtype to timezone-aware dtype is deprecated and will raise in a future version.  Use ser.dt.tz_localize instead.
  return col.astype(pandas_dtype)

```


Discussion:

Resolving patch:
diff --git a/pandera/engines/pandas_engine.py b/pandera/engines/pandas_engine.py
--- a/pandera/engines/pandas_engine.py
+++ b/pandera/engines/pandas_engine.py
@@ -730,31 +730,6 @@ def _get_to_datetime_fn(obj: Any) -> Callable:
 
         return to_datetime_fn
 
-    def _coerce(
-        self, data_container: PandasObject, pandas_dtype: Any
-    ) -> PandasObject:
-        to_datetime_fn = self._get_to_datetime_fn(data_container)
-
-        def _to_datetime(col: PandasObject) -> PandasObject:
-            col = to_datetime_fn(col, **self.to_datetime_kwargs)
-            return col.astype(pandas_dtype)
-
-        if isinstance(data_container, pd.DataFrame):
-            # pd.to_datetime transforms a df input into a series.
-            # We actually want to coerce every columns.
-            return data_container.transform(to_datetime_fn)
-
-        return _to_datetime(data_container)
-
-    def coerce(self, data_container: PandasObject) -> PandasObject:
-        return self._coerce(data_container, pandas_dtype=self.type)
-
-    def coerce_value(self, value: Any) -> Any:
-        """Coerce an value to specified datatime type."""
-        return self._get_to_datetime_fn(value)(
-            value, **self.to_datetime_kwargs
-        )
-
 
 @Engine.register_dtype(
     equivalents=[
@@ -777,13 +752,24 @@ class DateTime(_BaseDateTime, dtypes.Timestamp):
     )
     unit: str = "ns"
     """The precision of the datetime data. Currently limited to "ns"."""
+
     tz: Optional[datetime.tzinfo] = None
     """The timezone."""
+
     to_datetime_kwargs: Dict[str, Any] = dataclasses.field(
         default_factory=dict, compare=False, repr=False
     )
     "Any additional kwargs passed to :func:`pandas.to_datetime` for coercion."
 
+    tz_localize_kwargs: Dict[str, Any] = dataclasses.field(
+        default_factory=dict, compare=False, repr=False
+    )
+    "Keyword arguments passed to :func:`pandas.Series.dt.tz_localize` for coercion."
+
+    _default_tz_localize_kwargs = {
+        "ambiguous": "infer",
+    }
+
     def __post_init__(self):
         if self.tz is None:
             type_ = np.dtype("datetime64[ns]")
@@ -794,27 +780,35 @@ def __post_init__(self):
 
         object.__setattr__(self, "type", type_)
 
-    @staticmethod
-    def _get_to_datetime_fn(obj: Any) -> Callable:
-
-        # NOTE: this is a hack to support pyspark.pandas. This needs to be
-        # thoroughly tested, right now pyspark.pandas returns NA when a
-        # dtype value can't be coerced into the target dtype.
-        to_datetime_fn = pd.to_datetime
-        if type(obj).__module__.startswith(
-            "pyspark.pandas"
-        ):  # pragma: no cover
-            # pylint: disable=import-outside-toplevel
-            import pyspark.pandas as ps
+    def _coerce(
+        self, data_container: PandasObject, pandas_dtype: Any
+    ) -> PandasObject:
+        to_datetime_fn = self._get_to_datetime_fn(data_container)
+        _tz_localize_kwargs = {
+            **self._default_tz_localize_kwargs,
+            **self.tz_localize_kwargs,
+        }
 
-            to_datetime_fn = ps.to_datetime
-        if type(obj).__module__.startswith("modin.pandas"):
-            # pylint: disable=import-outside-toplevel
-            import modin.pandas as mpd
+        def _to_datetime(col: PandasObject) -> PandasObject:
+            col = to_datetime_fn(col, **self.to_datetime_kwargs)
+            if (
+                hasattr(pandas_dtype, "tz")
+                and pandas_dtype.tz is not None
+                and col.dt.tz is None
+            ):
+                # localize datetime column so that it's timezone-aware
+                col = col.dt.tz_localize(
+                    pandas_dtype.tz,
+                    **_tz_localize_kwargs,
+                )
+            return col.astype(pandas_dtype)
 
-            to_datetime_fn = mpd.to_datetime
+        if isinstance(data_container, pd.DataFrame):
+            # pd.to_datetime transforms a df input into a series.
+            # We actually want to coerce every columns.
+            return data_container.transform(_to_datetime)
 
-        return to_datetime_fn
+        return _to_datetime(data_container)
 
     @classmethod
     def from_parametrized_dtype(cls, pd_dtype: pd.DatetimeTZDtype):
@@ -822,6 +816,15 @@ def from_parametrized_dtype(cls, pd_dtype: pd.DatetimeTZDtype):
         a Pandera :class:`pandera.engines.pandas_engine.DateTime`."""
         return cls(unit=pd_dtype.unit, tz=pd_dtype.tz)  # type: ignore
 
+    def coerce(self, data_container: PandasObject) -> PandasObject:
+        return self._coerce(data_container, pandas_dtype=self.type)
+
+    def coerce_value(self, value: Any) -> Any:
+        """Coerce an value to specified datatime type."""
+        return self._get_to_datetime_fn(value)(
+            value, **self.to_datetime_kwargs
+        )
+
     def __str__(self) -> str:
         if self.type == np.dtype("datetime64[ns]"):
             return "datetime64[ns]"
@@ -849,17 +852,36 @@ class Date(_BaseDateTime, dtypes.Date):
 
     # define __init__ to please mypy
     def __init__(  # pylint:disable=super-init-not-called
-        self, to_datetime_kwargs: Optional[Dict[str, Any]] = None
+        self,
+        to_datetime_kwargs: Optional[Dict[str, Any]] = None,
     ) -> None:
         object.__setattr__(
             self, "to_datetime_kwargs", to_datetime_kwargs or {}
         )
 
+    def _coerce(
+        self, data_container: PandasObject, pandas_dtype: Any
+    ) -> PandasObject:
+        to_datetime_fn = self._get_to_datetime_fn(data_container)
+
+        def _to_datetime(col: PandasObject) -> PandasObject:
+            col = to_datetime_fn(col, **self.to_datetime_kwargs)
+            return col.astype(pandas_dtype).dt.date
+
+        if isinstance(data_container, pd.DataFrame):
+            # pd.to_datetime transforms a df input into a series.
+            # We actually want to coerce every columns.
+            return data_container.transform(_to_datetime)
+
+        return _to_datetime(data_container)
+
     def coerce(self, data_container: PandasObject) -> PandasObject:
-        return self._coerce(data_container, pandas_dtype=np.datetime64).dt.date
+        return self._coerce(data_container, pandas_dtype="datetime64[ns]")
 
     def coerce_value(self, value: Any) -> Any:
-        coerced = _BaseDateTime.coerce_value(self, value)
+        coerced = self._get_to_datetime_fn(value)(
+            value, **self.to_datetime_kwargs
+        )
         return coerced.date() if coerced is not None else pd.NaT
 
     def check(  # type: ignore
diff --git a/pandera/strategies.py b/pandera/strategies.py
--- a/pandera/strategies.py
+++ b/pandera/strategies.py
@@ -378,7 +378,7 @@ def compat_kwargs(*args):
         )
 
     np_dtype = to_numpy_dtype(pandera_dtype)
-    if strategy:
+    if strategy is not None:
         if _is_datetime_tz(pandera_dtype):
             return _datetime_strategy(pandera_dtype.type, strategy)  # type: ignore
         return strategy.map(np_dtype.type)
diff --git a/pandera/typing/__init__.py b/pandera/typing/__init__.py
--- a/pandera/typing/__init__.py
+++ b/pandera/typing/__init__.py
@@ -21,7 +21,9 @@
     AnnotationInfo,
     Bool,
     Category,
+    Date,
     DateTime,
+    Decimal,
     Float,
     Float16,
     Float32,


Test patch:
diff --git a/tests/core/test_dtypes.py b/tests/core/test_dtypes.py
--- a/tests/core/test_dtypes.py
+++ b/tests/core/test_dtypes.py
@@ -15,6 +15,7 @@
 from _pytest.mark.structures import ParameterSet
 from _pytest.python import Metafunc
 from hypothesis import strategies as st
+from pandas import DatetimeTZDtype
 
 import pandera as pa
 from pandera.engines import pandas_engine
@@ -388,20 +389,20 @@ def test_coerce_cast(dtypes, examples, data):
     assert expected_dtype.check(pandas_engine.Engine.dtype(coerced_dtype))
 
 
+dt_examples_to_coerce = [
+    "2021-09-01",
+    datetime.datetime(2021, 1, 7),
+    pd.NaT,
+    "not_a_date",
+]
+
+
 @pytest.mark.parametrize(
     "examples, type_, failure_indices",
     [
         (["a", 0, "b"], int, [0, 2]),
-        (
-            [
-                "2021-09-01",
-                datetime.datetime(2021, 1, 7),
-                pd.NaT,
-                "not_a_date",
-            ],
-            datetime.datetime,
-            [3],
-        ),
+        (dt_examples_to_coerce, datetime.datetime, [3]),
+        (dt_examples_to_coerce, DatetimeTZDtype(tz="UTC"), [3]),
     ],
 )
 def test_try_coerce(examples, type_, failure_indices):
diff --git a/tests/core/test_pandas_engine.py b/tests/core/test_pandas_engine.py
--- a/tests/core/test_pandas_engine.py
+++ b/tests/core/test_pandas_engine.py
@@ -1,8 +1,13 @@
-"""Test numpy engine."""
+"""Test pandas engine."""
 
+from datetime import date
+
+import hypothesis
+import hypothesis.extra.pandas as pd_st
 import hypothesis.strategies as st
 import pandas as pd
 import pytest
+import pytz
 from hypothesis import given
 
 from pandera.engines import pandas_engine
@@ -109,3 +114,80 @@ def test_pandas_boolean_native_type_error(data):
     for _, value in data.iteritems():
         with pytest.raises(TypeError):
             dtype.coerce_value(value)
+
+
+@hypothesis.settings(max_examples=1000)
+@pytest.mark.parametrize("timezone_aware", [True, False])
+@given(
+    data=pd_st.series(
+        dtype="datetime64[ns]",
+        index=pd_st.range_indexes(min_size=5, max_size=10),
+    ),
+    timezone=st.sampled_from(pytz.all_timezones),
+)
+def test_pandas_datetimetz_dtype(timezone_aware, data, timezone):
+    """
+    Test that pandas timezone-aware datetime correctly handles timezone-aware
+    and non-timezone-aware data.
+    """
+    timezone = pytz.timezone(timezone)
+    tz_localize_kwargs = {"ambiguous": "NaT"}
+
+    expected_failure = False
+    if timezone_aware:
+        data = data.dt.tz_localize(pytz.utc)
+    else:
+        assert data.dt.tz is None
+        try:
+            data.dt.tz_localize(timezone, **tz_localize_kwargs)
+        except pytz.exceptions.NonExistentTimeError:
+            expected_failure = True
+
+    # pylint: disable=unexpected-keyword-arg,no-value-for-parameter
+    dtype = pandas_engine.Engine.dtype(
+        pandas_engine.DateTime(
+            tz=timezone, tz_localize_kwargs=tz_localize_kwargs
+        )
+    )
+    if expected_failure:
+        with pytest.raises(pytz.exceptions.NonExistentTimeError):
+            dtype.coerce(data)
+    else:
+        coerced_data = dtype.coerce(data)
+        assert coerced_data.dt.tz == timezone
+
+
+@hypothesis.settings(max_examples=1000)
+@pytest.mark.parametrize("to_df", [True, False])
+@given(
+    data=pd_st.series(
+        dtype="datetime64[ns]",
+        index=pd_st.range_indexes(min_size=5, max_size=10),
+    ),
+)
+def test_pandas_date_coerce_dtype(to_df, data):
+    """Test that pandas Date dtype coerces to datetime.date object."""
+
+    data = data.to_frame() if to_df else data
+
+    dtype = pandas_engine.Engine.dtype(pandas_engine.Date())
+    coerced_data = dtype.coerce(data)
+
+    if to_df:
+        assert (coerced_data.dtypes == "object").all() or (
+            coerced_data.isna().all(axis=None)
+            and (coerced_data.dtypes == "datetime64[ns]").all()
+        )
+
+        assert (
+            coerced_data.applymap(lambda x: isinstance(x, date))
+            | coerced_data.isna()
+        ).all(axis=None)
+        return
+
+    assert (coerced_data.dtype == "object") or (
+        coerced_data.isna().all() and coerced_data.dtype == "datetime64[ns]"
+    )
+    assert (
+        coerced_data.map(lambda x: isinstance(x, date)) | coerced_data.isna()
+    ).all()