From 4df2138a2acb61582058fc07319168fb4986c44d Mon Sep 17 00:00:00 2001
From: Allen Porter <allen.porter@gmail.com>
Date: Thu, 6 Oct 2022 07:24:25 -0700
Subject: [PATCH] Add APIs for overriding the local timezone (#121)

* Add utility for overriding the local timezone

* Rename context manager to use_local_timezone
---
 ical/util.py           | 61 +++++++++++++++++++++++++++++++++++++-----
 tests/test_calendar.py | 46 +++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 7 deletions(-)

diff --git a/ical/util.py b/ical/util.py
index 1233e14..47f1146 100644
--- a/ical/util.py
+++ b/ical/util.py
@@ -4,15 +4,19 @@
 
 import datetime
 import uuid
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
 
 __all__ = [
-  "dtstamp_factory",
-  "uid_factory",
-  "local_timezone",
+    "use_local_timezone",
+    "dtstamp_factory",
+    "uid_factory",
 ]
 
 
 MIDNIGHT = datetime.time()
+LOCAL_TZ = ContextVar[datetime.tzinfo]("_local_tz")
 
 
 def dtstamp_factory() -> datetime.datetime:
@@ -25,12 +29,55 @@ def uid_factory() -> str:
     return str(uuid.uuid1())
 
 
+@contextmanager
+def use_local_timezone(local_tz: datetime.tzinfo) -> Generator[None, None, None]:
+    """Set the local timezone to use when converting a date to datetime.
+
+    This is expected to be used as a context manager when the default timezone
+    used by python is not the timezone to be used for calendar operations (the
+    attendees local timezone).
+
+    Example:
+    ```
+    import datetime
+    import zoneinfo
+    from ical.calendar import Calendar
+    from ical.event import Event
+    from ical.util import use_local_timezone
+
+    cal = Calendar()
+    cal.events.append(
+        Event(
+            summary="Example",
+            start=datetime.date(2022, 2, 1),
+            end=datetime.date(2022, 2, 2)
+        )
+    )
+    # Use UTC-8 as local timezone
+    with use_local_timezone(zoneinfo.ZoneInfo("America/Los_Angeles")):
+        # Returns event above
+        events = cal.timeline.start_after(
+            datetime.datetime(2022, 2, 1, 7, 59, 59, tzinfo=datetime.timezone.utc))
+
+        # Does not return event above
+        events = cal.timeline.start_after(
+            datetime.datetime(2022, 2, 1, 8, 00, 00, tzinfo=datetime.timezone.utc))
+    ```
+    """
+    orig_tz = LOCAL_TZ.set(local_tz)
+    try:
+        yield
+    finally:
+        LOCAL_TZ.reset(orig_tz)
+
+
 def local_timezone() -> datetime.tzinfo:
     """Get the local timezone to use when converting date to datetime."""
-    local_tz = datetime.datetime.now().astimezone().tzinfo
-    if not local_tz:
-        return datetime.timezone.utc
-    return local_tz
+    if local_tz := LOCAL_TZ.get(None):
+        return local_tz
+    if local_tz := datetime.datetime.now().astimezone().tzinfo:
+        return local_tz
+    return datetime.timezone.utc
 
 
 def normalize_datetime(value: datetime.date | datetime.datetime) -> datetime.datetime:
diff --git a/tests/test_calendar.py b/tests/test_calendar.py
index 5e6d7e0..520b052 100644
--- a/tests/test_calendar.py
+++ b/tests/test_calendar.py
@@ -4,6 +4,7 @@
 import datetime
 import re
 import uuid
+import zoneinfo
 from typing import Generator
 from unittest.mock import patch
 
@@ -13,6 +14,7 @@
 from ical.calendar import Calendar
 from ical.calendar_stream import IcsCalendarStream
 from ical.event import Event
+from ical.util import use_local_timezone
 
 
 @pytest.fixture(name="calendar")
@@ -275,3 +277,47 @@ def test_create_and_serialize_calendar(
         "END:VEVENT",
         "END:VCALENDAR",
     ]
+
+
+@pytest.mark.parametrize(
+    "tzname,dt_before,dt_after",
+    [
+        (
+            "America/Los_Angeles",  # UTC-8 in Feb
+            datetime.datetime(2000, 2, 1, 7, 59, 59, tzinfo=datetime.timezone.utc),
+            datetime.datetime(2000, 2, 1, 8, 0, 0, tzinfo=datetime.timezone.utc),
+        ),
+        (
+            "America/Regina",  # UTC-6 all year round
+            datetime.datetime(2000, 2, 1, 5, 59, 59, tzinfo=datetime.timezone.utc),
+            datetime.datetime(2000, 2, 1, 6, 0, 0, tzinfo=datetime.timezone.utc),
+        ),
+        (
+            "CET",  # UTC-1 in Feb
+            datetime.datetime(2000, 1, 31, 22, 59, 59, tzinfo=datetime.timezone.utc),
+            datetime.datetime(2000, 1, 31, 23, 0, 0, tzinfo=datetime.timezone.utc),
+        ),
+    ],
+)
+def test_all_day_with_local_timezone(
+    tzname: str, dt_before: datetime.datetime, dt_after: datetime.datetime
+) -> None:
+    """Test iteration of all day events using local timezone override."""
+    cal = Calendar()
+    cal.events.extend(
+        [
+            Event(
+                summary="event",
+                start=datetime.date(2000, 2, 1),
+                end=datetime.date(2000, 2, 2),
+            ),
+        ]
+    )
+
+    def start_after(dtstart: datetime.datetime) -> list[str]:
+        nonlocal cal
+        return [e.summary for e in cal.timeline.start_after(dtstart)]
+
+    with use_local_timezone(zoneinfo.ZoneInfo(tzname)):
+        assert start_after(dt_before) == ["event"]
+        assert not start_after(dt_after)