diff --git a/src/pendulum/tz/timezone.py b/src/pendulum/tz/timezone.py index e46e13df..1dd7ac3d 100644 --- a/src/pendulum/tz/timezone.py +++ b/src/pendulum/tz/timezone.py @@ -2,11 +2,13 @@ from __future__ import annotations import datetime as _datetime +import io import zoneinfo from abc import ABC from abc import abstractmethod from typing import TYPE_CHECKING +from typing import Any from typing import TypeVar from typing import cast @@ -16,6 +18,9 @@ if TYPE_CHECKING: + from collections.abc import Callable + from zoneinfo._common import _IOBytes + from typing_extensions import Self POST_TRANSITION = "post" @@ -60,12 +65,46 @@ class Timezone(zoneinfo.ZoneInfo, PendulumTimezone): >>> tz = Timezone('Europe/Paris') """ + _file_bytes: bytes | None = None + def __new__(cls, key: str) -> Self: try: return super().__new__(cls, key) # type: ignore[call-arg] except zoneinfo.ZoneInfoNotFoundError: raise InvalidTimezone(key) + @classmethod + def from_file(cls, fobj: _IOBytes, /, key: str | None = None) -> Self: + # The underlying zoneinfo.ZoneInfo.from_file() refuses to pickle any + # instance built this way, key or no key, since it has no record of + # which file it came from to reconstruct on unpickling. That's the + # path get_local_timezone() falls back to when the system's local + # timezone can't be identified by name (no /etc/timezone, no readable + # /etc/localtime symlink, etc.), so keep the raw TZif bytes around and + # use them to rebuild an equivalent instance if this ever needs to be + # pickled instead. + data = fobj.read(-1) + instance = cast("Self", super().from_file(io.BytesIO(data), key=key)) + instance._file_bytes = data + + return instance + + def __reduce__( + self, + ) -> ( + tuple[Callable[[bytes, str | None], Self], tuple[bytes, str | None]] + | str + | tuple[Any, ...] + ): + if self._file_bytes is None: + return super().__reduce__() + + return self.__class__._from_pickled_file, (self._file_bytes, self.key) + + @classmethod + def _from_pickled_file(cls, data: bytes, key: str | None) -> Self: + return cls.from_file(io.BytesIO(data), key=key) + def __eq__(self, other: object) -> bool: return isinstance(other, Timezone) and self.key == other.key diff --git a/tests/tz/test_timezone.py b/tests/tz/test_timezone.py index 3f090168..f4e41ec6 100644 --- a/tests/tz/test_timezone.py +++ b/tests/tz/test_timezone.py @@ -1,5 +1,6 @@ from __future__ import annotations +import pickle import zoneinfo from datetime import datetime @@ -462,3 +463,48 @@ def test_repr(): tz = timezone("Europe/Paris") assert repr(tz) == "Timezone('Europe/Paris')" + + +def _paris_tzif_bytes() -> bytes: + # tzdata bundles the same IANA zoneinfo files the system copy under + # /usr/share/zoneinfo would have, in a location that works on every + # platform pendulum supports (including Windows, which has no system + # zoneinfo directory at all). + from importlib import resources + + return resources.files("tzdata.zoneinfo").joinpath("Europe", "Paris").read_bytes() + + +def test_from_file_without_a_key_can_be_pickled(): + # get_local_timezone()'s last-resort fallback (no /etc/timezone, no + # readable /etc/localtime symlink to derive a name from) reads the raw + # zoneinfo file and builds a Timezone this same way, with no key. + from io import BytesIO + + import pendulum.tz.timezone as timezone_module + + tz = timezone_module.Timezone.from_file(BytesIO(_paris_tzif_bytes())) + + assert tz.key is None + + unpickled = pickle.loads(pickle.dumps(tz)) + + assert unpickled.key is None + dt = datetime(2024, 7, 1, 12, tzinfo=unpickled) + assert dt.utcoffset() == timedelta(hours=2) # CEST, matches Europe/Paris in July + dt = datetime(2024, 1, 1, 12, tzinfo=unpickled) + assert dt.utcoffset() == timedelta(hours=1) # CET, matches Europe/Paris in January + + +def test_from_file_with_a_key_still_pickles_by_key(): + from io import BytesIO + + import pendulum.tz.timezone as timezone_module + + tz = timezone_module.Timezone.from_file( + BytesIO(_paris_tzif_bytes()), key="Europe/Paris" + ) + + unpickled = pickle.loads(pickle.dumps(tz)) + + assert unpickled.key == "Europe/Paris"