Enhancement: Add support for timezone-flexible DateTime (unionai-oss#1352)

Signed-off-by: Max Raphael <[email protected]>

Author: max-raphael

pandera/engines/pandas_engine.py

@@ -862,7 +862,11 @@ class DateTime(_BaseDateTime, dtypes.Timestamp):
 
     timezone_flexible: bool = False
     """
+<<<<<<< HEAD
     A flag indicating whether the datetime data should be handled flexibly with respect to timezones. 
+=======
+    A flag indicating whether the datetime data should be handled flexibly with respect to timezones.
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
     When set to True, the function will ignore 'tz' and allow datetimes with any timezone(s). If coerce is set to True,
     the function can accept timezone-naive datetimes, and will convert all datetimes to the specified tz (or 'UTC').
     """
@@ -944,6 +948,7 @@ def from_parametrized_dtype(cls, pd_dtype: pd.DatetimeTZDtype):
 
     def coerce(self, data_container: PandasObject) -> PandasObject:
         if self.timezone_flexible:
+<<<<<<< HEAD
             data_container = self._prepare_coerce_timezone_flexible(data_container=data_container)
         return self._coerce(data_container, pandas_dtype=self.type)
 
@@ -953,18 +958,45 @@ def _prepare_coerce_timezone_flexible(self, data_container: PandasObject) -> Pan
             tz = self.tz if self.tz else data_container.dtype.tz
             unit = self.unit if self.unit else data_container.dtype.unit
             type_ = pd.DatetimeTZDtype(unit, tz)
+=======
+            data_container = self._prepare_coerce_timezone_flexible(
+                data_container=data_container
+            )
+        return self._coerce(data_container, pandas_dtype=self.type)
+
+    def _prepare_coerce_timezone_flexible(
+        self, data_container: PandasObject
+    ) -> PandasObject:
+        # If there is a single timezone, define the type as a timezone-aware DatetimeTZDtype
+        if isinstance(data_container.dtype, pd.DatetimeTZDtype):
+            tz = self.tz if self.tz else data_container.dtype.tz
+            type_ = pd.DatetimeTZDtype("ns", tz)
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
             object.__setattr__(self, "tz", tz)
             object.__setattr__(self, "type", type_)
         # If there are multiple timezones, convert them to the specified tz (default 'UTC') and set the type accordingly
         elif all(isinstance(x, datetime.datetime) for x in data_container):
             container_type = type(data_container)
+<<<<<<< HEAD
             tz = self.tz if self.tz else 'UTC'
             unit = self.unit if self.unit else data_container.dtype.unit
             data_container = container_type(
                 [pd.Timestamp(ts).tz_convert(tz) if pd.Timestamp(ts).tzinfo else pd.Timestamp(ts).tz_localize(tz)
                  for ts in data_container]
             )
             type_ = pd.DatetimeTZDtype(unit, tz)
+=======
+            tz = self.tz if self.tz else "UTC"  # type: ignore
+            data_container = container_type(
+                [
+                    pd.Timestamp(ts).tz_convert(tz)  # type: ignore
+                    if pd.Timestamp(ts).tzinfo  # type: ignore
+                    else pd.Timestamp(ts).tz_localize(tz)  # type: ignore
+                    for ts in data_container
+                ]
+            )
+            type_ = pd.DatetimeTZDtype("ns", tz)
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
             object.__setattr__(self, "tz", tz)
             object.__setattr__(self, "type", type_)
         else:
@@ -984,6 +1016,7 @@ def check(
         data_container: Optional[PandasObject] = None,
     ) -> Union[bool, Iterable[bool]]:
         if self.timezone_flexible:
+<<<<<<< HEAD
             self._prepare_check_timezone_flexible(pandera_dtype=pandera_dtype, data_container=data_container)
         return super().check(pandera_dtype, data_container)
 
@@ -998,6 +1031,34 @@ def _prepare_check_timezone_flexible(
         # If the data has a mix of timezones, pandas defines the dtype as 'object
         elif all(isinstance(x, datetime.datetime) and x.tzinfo is not None for x in data_container):
             object.__setattr__(self, "type", np.dtype('O'))
+=======
+            self._prepare_check_timezone_flexible(
+                pandera_dtype=pandera_dtype, data_container=data_container
+            )
+        return super().check(pandera_dtype, data_container)
+
+    def _prepare_check_timezone_flexible(
+        self,
+        pandera_dtype: dtypes.DataType,
+        data_container: Optional[PandasObject],
+    ) -> None:
+        if data_container is None:
+            return
+        # If there is a single timezone, define the type as a timezone-aware DatetimeTZDtype
+        if (
+            isinstance(pandera_dtype, DateTime)
+            and pandera_dtype.tz is not None
+        ):
+            type_ = pd.DatetimeTZDtype("ns", pandera_dtype.tz)
+            object.__setattr__(self, "tz", pandera_dtype.tz)
+            object.__setattr__(self, "type", type_)
+        # If the data has a mix of timezones, pandas defines the dtype as 'object
+        elif all(
+            isinstance(x, datetime.datetime) and x.tzinfo is not None
+            for x in data_container
+        ):
+            object.__setattr__(self, "type", np.dtype("O"))
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
         else:
             # Prepare to raise exception, adding type strictly for the check_dtype error message
             object.__setattr__(self, "type", "datetime64[ns, <timezone>]")

tests/core/test_pandas_engine.py

@@ -1,5 +1,6 @@
 """Test pandas engine."""
 
+<<<<<<< HEAD
 <<<<<<< HEAD
 import datetime as dt
 from typing import Tuple, List, Optional
@@ -8,6 +9,11 @@
 from datetime import date
 from typing import Any, Set
 >>>>>>> add pandas pyarrow backend support  (#1628)
+=======
+import datetime as dt
+from typing import Any, Set, Tuple, List, Optional
+from zoneinfo import ZoneInfo
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
 
 import hypothesis
 import hypothesis.extra.pandas as pd_st
@@ -210,7 +216,13 @@ def test_pandas_datetimetz_dtype(timezone_aware, data, timezone):
 
 
 def generate_test_cases_timezone_flexible() -> List[
+<<<<<<< HEAD
     Tuple[List[dt.datetime], Optional[dt.tzinfo], bool, List[dt.datetime], bool]
+=======
+    Tuple[
+        List[dt.datetime], Optional[dt.tzinfo], bool, List[dt.datetime], bool
+    ]
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
 ]:
     """
     Generate test parameter combinations for a given list of datetime lists.
@@ -226,6 +238,7 @@ def generate_test_cases_timezone_flexible() -> List[
     datetimes = [
         # multi tz and tz naive
         [
+<<<<<<< HEAD
             dt.datetime(2023, 3, 1, 4, tzinfo=ZoneInfo('America/New_York')),
             dt.datetime(2023, 3, 1, 5, tzinfo=ZoneInfo('America/Los_Angeles')),
             dt.datetime(2023, 3, 1, 5)
@@ -245,15 +258,45 @@ def generate_test_cases_timezone_flexible() -> List[
             dt.datetime(2023, 3, 1, 4, tzinfo=ZoneInfo('America/New_York')),
             dt.datetime(2023, 3, 1, 5, tzinfo=ZoneInfo('America/New_York'))
         ]
+=======
+            dt.datetime(2023, 3, 1, 4, tzinfo=ZoneInfo("America/New_York")),
+            dt.datetime(2023, 3, 1, 5, tzinfo=ZoneInfo("America/Los_Angeles")),
+            dt.datetime(2023, 3, 1, 5),
+        ],
+        # multiz tz
+        [
+            dt.datetime(2023, 3, 1, 4, tzinfo=ZoneInfo("America/New_York")),
+            dt.datetime(2023, 3, 1, 5, tzinfo=ZoneInfo("America/Los_Angeles")),
+        ],
+        # tz naive
+        [dt.datetime(2023, 3, 1, 4), dt.datetime(2023, 3, 1, 5)],
+        # single tz
+        [
+            dt.datetime(2023, 3, 1, 4, tzinfo=ZoneInfo("America/New_York")),
+            dt.datetime(2023, 3, 1, 5, tzinfo=ZoneInfo("America/New_York")),
+        ],
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
     ]
 
     test_cases = []
 
     for datetime_list in datetimes:
         for coerce in [True, False]:
+<<<<<<< HEAD
             for tz in [None, ZoneInfo("America/Chicago"), dt.timezone(dt.timedelta(hours=2))]:
                 # Determine if the test should raise an exception
                 has_naive_datetime = any([dt.tzinfo is None for dt in datetime_list])
+=======
+            for tz in [
+                None,
+                ZoneInfo("America/Chicago"),
+                dt.timezone(dt.timedelta(hours=2)),
+            ]:
+                # Determine if the test should raise an exception
+                has_naive_datetime = any(
+                    dt.tzinfo is None for dt in datetime_list
+                )
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
                 raises = has_naive_datetime and not coerce
 
                 # Generate expected output
@@ -264,52 +307,112 @@ def generate_test_cases_timezone_flexible() -> List[
                         # localize / convert the input datetimes to the specified tz or 'UTC' (default)
                         use_tz = tz if tz else ZoneInfo("UTC")
                         expected_output_naive = [
+<<<<<<< HEAD
                             dt.replace(tzinfo=use_tz) for dt in datetime_list if dt.tzinfo is None
                         ]
                         expected_output_aware = [
                             dt.astimezone(use_tz) for dt in datetime_list if dt.tzinfo is not None
                         ]
                         expected_output = expected_output_naive + expected_output_aware
+=======
+                            dt.replace(tzinfo=use_tz)
+                            for dt in datetime_list
+                            if dt.tzinfo is None
+                        ]
+                        expected_output_aware = [
+                            dt.astimezone(use_tz)
+                            for dt in datetime_list
+                            if dt.tzinfo is not None
+                        ]
+                        expected_output = (
+                            expected_output_naive + expected_output_aware
+                        )
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
                     else:
                         # ignore tz
                         expected_output = datetime_list
 
+<<<<<<< HEAD
                 test_case = (datetime_list, tz, coerce, expected_output, raises)
                 test_cases.append(test_case)
 
     # define final test cases with improper type
     datetime_list = [dt.datetime(2023, 3, 1, 4, tzinfo=ZoneInfo('America/New_York')), "hello world"]
+=======
+                test_case = (
+                    datetime_list,
+                    tz,
+                    coerce,
+                    expected_output,
+                    raises,
+                )
+                test_cases.append(test_case)
+
+    # define final test cases with improper type
+    datetime_list = [
+        dt.datetime(2023, 3, 1, 4, tzinfo=ZoneInfo("America/New_York")),
+        "hello world",
+    ]
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
     tz = None
     expected_output = None
     raises = True
 
     bad_type_coerce = (datetime_list, tz, True, expected_output, raises)
     bad_type_no_coerce = (datetime_list, tz, False, expected_output, raises)
+<<<<<<< HEAD
     test_cases.extend([bad_type_coerce, bad_type_no_coerce])
 
     return test_cases
+=======
+    test_cases.extend([bad_type_coerce, bad_type_no_coerce])  # type: ignore
+
+    return test_cases  # type: ignore
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
 
 
 @pytest.mark.parametrize(
     "examples, tz, coerce, expected_output, raises",
+<<<<<<< HEAD
     generate_test_cases_timezone_flexible()
+=======
+    generate_test_cases_timezone_flexible(),
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
 )
 def test_dt_timezone_flexible(examples, tz, coerce, expected_output, raises):
     """Test that timezone_flexible works as expected"""
 
     # Testing using a pandera DataFrameModel rather than directly calling dtype coerce or validate because with
     # timezone_flexible, dtype is set dynamically based on the input data
     class SimpleSchema(DataFrameModel):
+<<<<<<< HEAD
         datetime_column: pandas_engine.DateTime(timezone_flexible=True, tz=tz) = Field(coerce=coerce)
 
     data = pd.DataFrame({'datetime_column': examples})
+=======
+        """Simple DF Model for testing"""
+
+        datetime_column: pandas_engine.DateTime(  # pylint: disable=unexpected-keyword-arg, no-value-for-parameter
+            timezone_flexible=True, tz=tz
+        ) = Field(
+            coerce=coerce
+        )
+
+    data = pd.DataFrame({"datetime_column": examples})
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
 
     if raises:
         with pytest.raises(SchemaError):
             SimpleSchema.validate(data)
     else:
         validated_df = SimpleSchema.validate(data)
+<<<<<<< HEAD
         assert sorted(validated_df['datetime_column'].tolist()) == sorted(expected_output)
+=======
+        assert sorted(validated_df["datetime_column"].tolist()) == sorted(
+            expected_output
+        )
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
 
 
 @hypothesis.settings(max_examples=1000)
@@ -344,7 +447,12 @@ def test_pandas_date_coerce_dtype(to_df, data):
         coerced_data.isna().all() and coerced_data.dtype == "datetime64[ns]"
     )
     assert (
+<<<<<<< HEAD
         coerced_data.map(lambda x: isinstance(x, dt.date)) | coerced_data.isna()
+=======
+        coerced_data.map(lambda x: isinstance(x, dt.date))
+        | coerced_data.isna()
+>>>>>>> Enhancement: Add support for timezone-flexible DateTime (#1352)
     ).all()
 
 
@@ -356,8 +464,8 @@ def test_pandas_date_coerce_dtype(to_df, data):
         pyarrow.struct([("foo", pyarrow.int64()), ("bar", pyarrow.string())]),
     ),
     (pd.Series([None, pd.NA, np.nan]), pyarrow.null),
-    (pd.Series([None, date(1970, 1, 1)]), pyarrow.date32),
-    (pd.Series([None, date(1970, 1, 1)]), pyarrow.date64),
+    (pd.Series([None, dt.date(1970, 1, 1)]), pyarrow.date32),
+    (pd.Series([None, dt.date(1970, 1, 1)]), pyarrow.date64),
     (pd.Series([1, 2]), pyarrow.duration("ns")),
     (pd.Series([1, 1e3, 1e6, 1e9, None]), pyarrow.time32("ms")),
     (pd.Series([1, 1e3, 1e6, 1e9, None]), pyarrow.time64("ns")),
@@ -402,8 +510,8 @@ def test_pandas_arrow_dtype(data, dtype):
         pyarrow.struct([("foo", pyarrow.string()), ("bar", pyarrow.int64())]),
     ),
     (pd.Series(["a", "1"]), pyarrow.null),
-    (pd.Series(["a", date(1970, 1, 1), "1970-01-01"]), pyarrow.date32),
-    (pd.Series(["a", date(1970, 1, 1), "1970-01-01"]), pyarrow.date64),
+    (pd.Series(["a", dt.date(1970, 1, 1), "1970-01-01"]), pyarrow.date32),
+    (pd.Series(["a", dt.date(1970, 1, 1), "1970-01-01"]), pyarrow.date64),
     (pd.Series(["a"]), pyarrow.duration("ns")),
     (pd.Series(["a", "b"]), pyarrow.time32("ms")),
     (pd.Series(["a", "b"]), pyarrow.time64("ns")),

tests/pyspark/test_schemas_on_pyspark_pandas.py

@@ -292,6 +292,7 @@ def test_nullable(
     data: st.DataObject,
 ):
     """Test nullable checks on pyspark.pandas dataframes."""
+    print(dtype)
 
     if version.parse(np.__version__) >= version.parse(
         "1.24.0"