fixup! feat: make temporal partition_mapper timezone aware

Author: Lee-W

airflow-core/src/airflow/partition_mappers/temporal.py

@@ -18,13 +18,14 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from typing import Any
-
-from pendulum import FixedTimezone, Timezone
+from typing import TYPE_CHECKING, Any
 
 from airflow._shared.timezones.timezone import parse_timezone
 from airflow.partition_mappers.base import PartitionMapper
 
+if TYPE_CHECKING:
+    from pendulum import FixedTimezone, Timezone
+
 
 class _BaseTemporalMapper(PartitionMapper, ABC):
     """Base class for Temporal Partition Mappers."""
@@ -33,10 +34,10 @@ class _BaseTemporalMapper(PartitionMapper, ABC):
 
     def __init__(
         self,
-        *
-        input_format: str = "%Y-%m-%dT%H:%M:%S",
-        output_format: str | None = None,
+        *,
         timezone: str | Timezone | FixedTimezone,
+        input_format: str = "%Y-%m-%dT%H:%M:%S%z",
+        output_format: str | None = None,
     ):
         self.input_format = input_format
         self.output_format = output_format or self.default_output_format
@@ -58,14 +59,18 @@ def format(self, dt: datetime) -> str:
         return dt.strftime(self.output_format)
 
     def serialize(self) -> dict[str, Any]:
+        from airflow.serialization.encoders import encode_timezone
+
         return {
+            "timezone": encode_timezone(self._timezone),
             "input_format": self.input_format,
             "output_format": self.output_format,
         }
 
     @classmethod
     def deserialize(cls, data: dict[str, Any]) -> PartitionMapper:
         return cls(
+            timezone=parse_timezone(data["timezone"]),
             input_format=data["input_format"],
             output_format=data["output_format"],
         )
@@ -74,7 +79,7 @@ def deserialize(cls, data: dict[str, Any]) -> PartitionMapper:
 class HourlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to hour."""
 
-    default_output_format = "%Y-%m-%dT%H"
+    default_output_format = "%Y-%m-%dT%H%z"
 
     def normalize(self, dt: datetime) -> datetime:
         return dt.replace(minute=0, second=0, microsecond=0)
@@ -83,7 +88,7 @@ def normalize(self, dt: datetime) -> datetime:
 class DailyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to day."""
 
-    default_output_format = "%Y-%m-%d"
+    default_output_format = "%Y-%m-%d%z"
 
     def normalize(self, dt: datetime) -> datetime:
         return dt.replace(hour=0, minute=0, second=0, microsecond=0)
@@ -92,7 +97,7 @@ def normalize(self, dt: datetime) -> datetime:
 class WeeklyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to week."""
 
-    default_output_format = "%Y-%m-%d (W%V)"
+    default_output_format = "%Y-%m-%d (W%V)%z"
 
     def normalize(self, dt: datetime) -> datetime:
         start = dt - timedelta(days=dt.weekday())
@@ -102,7 +107,7 @@ def normalize(self, dt: datetime) -> datetime:
 class MonthlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to month."""
 
-    default_output_format = "%Y-%m"
+    default_output_format = "%Y-%m%z"
 
     def normalize(self, dt: datetime) -> datetime:
         return dt.replace(
@@ -117,7 +122,7 @@ def normalize(self, dt: datetime) -> datetime:
 class QuarterlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to quarter."""
 
-    default_output_format = "%Y-Q{quarter}"
+    default_output_format = "%Y-Q{quarter}%z"
 
     def normalize(self, dt: datetime) -> datetime:
         quarter = (dt.month - 1) // 3
@@ -139,7 +144,7 @@ def format(self, dt: datetime) -> str:
 class YearlyMapper(_BaseTemporalMapper):
     """Map a time-based partition key to year."""
 
-    default_output_format = "%Y"
+    default_output_format = "%Y%z"
 
     def normalize(self, dt: datetime) -> datetime:
         return dt.replace(