coerce dt indexes and series

Signed-off-by: Cristian Matache <[email protected]>

pandera/engines/pandas_engine.py

@@ -808,16 +808,23 @@ def _coerce(
 
         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,
-                )
+            if hasattr(pandas_dtype, "tz") and pandas_dtype.tz is not None:
+                if hasattr(col, "dt") 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,
+                    )
+                elif (
+                    hasattr(col, "tz")
+                    and col.tz != pandas_dtype.tz
+                    and hasattr(col, "tz_localize")
+                ):
+                    # localize datetime index sso that it's timezone-aware
+                    col = col.tz_localize(
+                        pandas_dtype.tz,
+                        **_tz_localize_kwargs,
+                    )
             return col.astype(pandas_dtype)
 
         if isinstance(data_container, pd.DataFrame):

tests/core/test_dtypes.py

@@ -15,7 +15,7 @@
 from _pytest.mark.structures import ParameterSet
 from _pytest.python import Metafunc
 from hypothesis import strategies as st
-from pandas import DatetimeTZDtype
+from pandas import DatetimeTZDtype, to_datetime
 
 import pandera as pa
 from pandera.engines import pandas_engine
@@ -419,6 +419,49 @@ def test_try_coerce(examples, type_, failure_indices):
         assert exc.failure_cases["index"].to_list() == failure_indices
 
 
+@pytest.mark.parametrize(
+    "examples, type_, has_tz",
+    [
+        (
+            ["2022-04-30T00:00:00Z", "2022-04-30T00:00:01Z"],
+            DatetimeTZDtype(tz="UTC"),
+            True,
+        ),
+        (
+            ["2022-04-30T00:00:00", "2022-04-30T00:00:01"],
+            DatetimeTZDtype(tz="UTC"),
+            True,
+        ),
+        (
+            ["2022-04-30T00:00:00Z", "2022-04-30T00:00:01Z"],
+            pandas_engine.DateTime,
+            False,
+        ),
+        (
+            ["2022-04-30T00:00:00", "2022-04-30T00:00:01"],
+            pandas_engine.DateTime,
+            False,
+        ),
+    ],
+)
+@pytest.mark.parametrize("is_index", [True, False])
+def test_coerce_dt(examples, type_, has_tz, is_index):
+    """Test coercion of Series and Indexes to DateTime with and without Time Zones."""
+    data = pd.Index(examples) if is_index else pd.Series(examples)
+    data_type = pandas_engine.Engine.dtype(type_)
+    if has_tz:
+        expected = [
+            to_datetime("2022-04-30 00:00:00", utc=True),
+            to_datetime("2022-04-30 00:00:01", utc=True),
+        ]
+    else:
+        expected = [
+            to_datetime("2022-04-30 00:00:00"),
+            to_datetime("2022-04-30 00:00:01"),
+        ]
+    assert data_type.try_coerce(data).to_list() == expected
+
+
 def test_coerce_string():
     """Test that strings can be coerced."""
     data = pd.Series([1, None], dtype="Int32")