Fixed TZ with DST on calendar operations

temporian/core/operators/calendar/base.py

@@ -17,6 +17,7 @@
 import datetime
 from typing import Union
 from abc import ABC, abstractmethod
+from dateutil.tz.tz import tzoffset
 
 from temporian.core.data.dtype import DType
 from temporian.core.data.node import (
@@ -27,29 +28,29 @@
 from temporian.proto import core_pb2 as pb
 
 
-def timezone_to_utc_offset(timezone: Union[str, float, int]) -> float:
-    """Normalizes timezone (tz name or number) to a float (UTC offset)."""
-    if isinstance(timezone, (bytes, str)):
-        tz_time = datetime.datetime.now(tz=pytz.timezone(timezone))
-        offset_timedelta = datetime.datetime.utcoffset(tz_time)
-        return (
-            0.0
-            if offset_timedelta is None
-            else offset_timedelta.total_seconds() / 3600
-        )
-    if not isinstance(timezone, (int, float)):
+def parse_tz(tz: Union[str, float, int]) -> datetime.tzinfo:
+    """Normalizes timezone (tz name or number) to a tzinfo."""
+    if isinstance(tz, (bytes, str)):
+        return pytz.timezone(tz)
+    if isinstance(tz, (int, float)):
+        if abs(tz) >= 24:
+            raise ValueError(
+                "Timezone offset must be a number of hours strictly between"
+                f" (-24, 24). Got: {tz=}"
+            )
+        return tzoffset(name=None, offset=datetime.timedelta(hours=tz))
+    else:
         raise TypeError(
             "Timezone argument (tz) must be a number of hours (int or"
             " float) between (-24, 24) or a timezone name (string, see"
-            f" pytz.all_timezones). Got '{timezone}' ({type(timezone)})."
+            f" pytz.all_timezones). Got '{tz}' ({type(tz)})."
         )
-    return float(timezone)
 
 
 class BaseCalendarOperator(Operator, ABC):
     """Interface definition and common logic for calendar operators."""
 
-    def __init__(self, sampling: EventSetNode, utc_offset: float):
+    def __init__(self, sampling: EventSetNode, tz: Union[int, float, str]):
         super().__init__()
 
         if not sampling.schema.is_unix_timestamp:
@@ -59,14 +60,8 @@ def __init__(self, sampling: EventSetNode, utc_offset: float):
                 " `is_unix_timestamp=True` when manually creating a sampling."
             )
 
-        # attribute: timezone
-        if abs(utc_offset) >= 24:
-            raise ValueError(
-                "Timezone offset must be a number of hours strictly between"
-                f" (-24, 24). Got: {utc_offset=}"
-            )
-        self._utc_offset = utc_offset
-        self.add_attribute("utc_offset", utc_offset)
+        self._tz = parse_tz(tz)
+        self.add_attribute("tz", tz)
 
         # input and output
         self.add_input("sampling", sampling)
@@ -87,17 +82,17 @@ def build_op_definition(cls) -> pb.OperatorDef:
             inputs=[pb.OperatorDef.Input(key="sampling")],
             attributes=[
                 pb.OperatorDef.Attribute(
-                    key="utc_offset",
-                    type=pb.OperatorDef.Attribute.Type.FLOAT_64,
+                    key="tz",
+                    type=pb.OperatorDef.Attribute.Type.ANY,
                 ),
             ],
             outputs=[pb.OperatorDef.Output(key="output")],
         )
 
     @property
-    def utc_offset(self) -> float:
+    def tz(self) -> datetime.tzinfo:
         """Gets timezone offset from UTC, in hours."""
-        return self._utc_offset
+        return self._tz
 
     @classmethod
     @abstractmethod

temporian/core/operators/calendar/day_of_month.py

@@ -14,13 +14,11 @@
 
 """Calendar day of month operator class and public API function definitions."""
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -42,6 +40,5 @@ def calendar_day_of_month(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarDayOfMonthOperator(sampling, utc_offset).outputs["output"]
+    return CalendarDayOfMonthOperator(sampling, tz).outputs["output"]

temporian/core/operators/calendar/day_of_week.py

@@ -15,13 +15,11 @@
 """Calendar day of week operator class and public API function definitions."""
 
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -43,6 +41,5 @@ def calendar_day_of_week(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarDayOfWeekOperator(sampling, utc_offset).outputs["output"]
+    return CalendarDayOfWeekOperator(sampling, tz).outputs["output"]

temporian/core/operators/calendar/day_of_year.py

@@ -15,13 +15,11 @@
 """Calendar day of year operator class and public API function definitions."""
 
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -43,6 +41,5 @@ def calendar_day_of_year(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarDayOfYearOperator(sampling, utc_offset).outputs["output"]
+    return CalendarDayOfYearOperator(sampling, tz).outputs["output"]

temporian/core/operators/calendar/hour.py

@@ -15,13 +15,11 @@
 """Calendar hour operator class and public API function definitions."""
 
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -43,6 +41,5 @@ def calendar_hour(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarHourOperator(sampling, utc_offset).outputs["output"]
+    return CalendarHourOperator(sampling, tz).outputs["output"]

temporian/core/operators/calendar/iso_week.py

@@ -15,13 +15,11 @@
 """Calendar ISO week operator class and public API function definitions."""
 
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -43,6 +41,5 @@ def calendar_iso_week(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarISOWeekOperator(sampling, utc_offset).outputs["output"]
+    return CalendarISOWeekOperator(sampling, tz).outputs["output"]

temporian/core/operators/calendar/minute.py

@@ -15,13 +15,11 @@
 """Calendar minute operator class and public API function definitions."""
 
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -43,6 +41,5 @@ def calendar_minute(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarMinuteOperator(sampling, utc_offset).outputs["output"]
+    return CalendarMinuteOperator(sampling, tz).outputs["output"]

temporian/core/operators/calendar/month.py

@@ -15,13 +15,11 @@
 """Calendar month operator class and public API function definitions."""
 
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -43,6 +41,5 @@ def calendar_month(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarMonthOperator(sampling, utc_offset).outputs["output"]
+    return CalendarMonthOperator(sampling, tz).outputs["output"]

temporian/core/operators/calendar/second.py

@@ -15,13 +15,11 @@
 """Calendar second operator class and public API function definitions."""
 
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -43,6 +41,5 @@ def calendar_second(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarSecondOperator(sampling, utc_offset).outputs["output"]
+    return CalendarSecondOperator(sampling, tz).outputs["output"]

temporian/core/operators/calendar/test/test_calendar_hour.py

@@ -41,6 +41,24 @@ def test_basic(self):
         result = evset.calendar_hour()
         assertOperatorResult(self, result, expected)
 
+    def test_dst(self):
+        timestamps = [
+            "2023-03-12 00:00:00",  # before DST -8
+            "2023-03-13 00:00:00",  # after DST -7
+        ]
+        evset = event_set(timestamps=timestamps)
+
+        expected = event_set(
+            timestamps=timestamps,
+            features={
+                "calendar_hour": i32([16, 17]),
+            },
+            same_sampling_as=evset,
+        )
+
+        result = evset.calendar_hour(tz="US/Pacific")
+        assertOperatorResult(self, result, expected)
+
 
 if __name__ == "__main__":
     absltest.main()

temporian/core/operators/calendar/year.py

@@ -15,13 +15,11 @@
 """Calendar year operator class and public API function definitions."""
 
 from typing import Union
+
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import EventSetNode
-from temporian.core.operators.calendar.base import (
-    BaseCalendarOperator,
-    timezone_to_utc_offset,
-)
+from temporian.core.operators.calendar.base import BaseCalendarOperator
 from temporian.core.typing import EventSetOrNode
 
 
@@ -43,6 +41,5 @@ def calendar_year(
     sampling: EventSetOrNode, tz: Union[str, float, int] = 0
 ) -> EventSetOrNode:
     assert isinstance(sampling, EventSetNode)
-    utc_offset = timezone_to_utc_offset(tz)
 
-    return CalendarYearOperator(sampling, utc_offset).outputs["output"]
+    return CalendarYearOperator(sampling, tz).outputs["output"]

temporian/implementation/numpy/operators/calendar/base.py

@@ -34,15 +34,14 @@ def __init__(self, operator: BaseCalendarOperator) -> None:
     def __call__(self, sampling: EventSet) -> Dict[str, EventSet]:
         assert isinstance(self.operator, BaseCalendarOperator)
         output_schema = self.output_schema("output")
-        tzinfo = timezone(timedelta(hours=self.operator.utc_offset))
 
         # create destination EventSet
         dst_evset = EventSet(data={}, schema=output_schema)
         for index_key, index_data in sampling.data.items():
             value = np.array(
                 [
                     self._get_value_from_datetime(
-                        datetime.fromtimestamp(ts, tz=tzinfo)
+                        datetime.fromtimestamp(ts, tz=self.operator.tz)
                     )
                     for ts in index_data.timestamps
                 ],