Handle localized timestamps in multiindex (#1426)

Signed-off-by: Rob <[email protected]>

pandera/strategies/pandas_strategies.py

@@ -259,11 +259,8 @@ def convert_dtype(array: Union[pd.Series, pd.Index], col_dtype: Any):
     return array.astype(col_dtype)
 
 
-def convert_dtypes(df: pd.DataFrame, col_dtypes: Dict[str, Any]):
+def convert_dtypes(df: pd.DataFrame, col_dtypes: Dict[Union[int, str], Any]):
     """Convert datatypes of a dataframe."""
-    if df.empty:
-        return df
-
     for col_name, col_dtype in col_dtypes.items():
         array: pd.Series = df[col_name]  # type: ignore[assignment]
         df[col_name] = convert_dtype(array, col_dtype)
@@ -1237,9 +1234,7 @@ def multiindex_strategy(
         index=pdst.range_indexes(
             min_size=0 if size is None else size, max_size=size
         ),
-    ).map(
-        lambda x: x.astype(index_dtypes)  # type: ignore[arg-type]
-    )
+    ).map(partial(convert_dtypes, col_dtypes=index_dtypes))
 
     # this is a hack to convert np.str_ data values into native python str.
     for name, dtype in index_dtypes.items():

tests/strategies/test_strategies.py

@@ -485,6 +485,15 @@ def test_dataframe_strategy(data_type, size, data):
         )
 
 
+def test_dataframe_strategy_empty_localized():
+    """Test DataFrameSchema strategy localizes timezones when empty."""
+    schema = pa.DataFrameSchema(
+        {"localized": pa.Column(pd.DatetimeTZDtype(tz="UTC", unit="ns"))}
+    )
+    example = schema.example(0)
+    schema(example)
+
+
 @pytest.mark.parametrize("size", [None, 0, 1, 3, 5])
 @hypothesis.given(st.data())
 def test_dataframe_strategy_with_check(size, data):
@@ -647,6 +656,7 @@ def test_multiindex_example() -> None:
             pa.Index(data_type, unique=True, name="level_0"),
             pa.Index(data_type, nullable=True),
             pa.Index(data_type),
+            pa.Index(pd.DatetimeTZDtype(tz="UTC", unit="ns")),
         ]
     )
     for _ in range(10):