Add time travel

Author: sdispater

pendulum/__init__.py

@@ -39,6 +39,7 @@
 from pendulum.helpers import week_starts_at
 from pendulum.parser import parse
 from pendulum.period import Period
+from pendulum.testing.traveller import Traveller
 from pendulum.time import Time
 from pendulum.tz import UTC
 from pendulum.tz import local_timezone
@@ -297,6 +298,15 @@ def period(start: DateTime, end: DateTime, absolute: bool = False) -> Period:
     return Period(start, end, absolute=absolute)
 
 
+# Testing
+
+_traveller = Traveller(DateTime)
+
+freeze = _traveller.freeze
+travel = _traveller.travel
+travel_to = _traveller.travel_to
+travel_back = _traveller.travel_back
+
 __all__ = [
     "__version__",
     "DAYS_PER_WEEK",
@@ -324,6 +334,7 @@ def period(start: DateTime, end: DateTime, absolute: bool = False) -> Period:
     "datetime",
     "duration",
     "format_diff",
+    "freeze",
     "from_format",
     "from_timestamp",
     "get_locale",
@@ -352,6 +363,9 @@ def period(start: DateTime, end: DateTime, absolute: bool = False) -> Period:
     "timezones",
     "today",
     "tomorrow",
+    "travel",
+    "travel_back",
+    "travel_to",
     "FixedTimezone",
     "Timezone",
     "yesterday",

pendulum/testing/__init__.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import cast
+
+from pendulum.datetime import DateTime
+from pendulum.utils._compat import PYPY
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+
+class BaseTraveller:
+    def __init__(self, datetime_class: type[DateTime] = DateTime) -> None:
+        self._datetime_class: type[DateTime] = datetime_class
+
+    def freeze(self: BaseTraveller) -> BaseTraveller:
+        raise NotImplementedError()
+
+    def travel_back(self: BaseTraveller) -> BaseTraveller:
+        raise NotImplementedError()
+
+    def travel(
+        self,
+        years: int = 0,
+        months: int = 0,
+        weeks: int = 0,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: int = 0,
+        microseconds: int = 0,
+    ) -> BaseTraveller:
+        raise NotImplementedError()
+
+    def travel_to(self, dt: DateTime) -> BaseTraveller:
+        raise NotImplementedError()
+
+
+if not PYPY:
+    import time_machine
+
+    class Traveller(BaseTraveller):
+        def __init__(self, datetime_class: type[DateTime] = DateTime) -> None:
+            super().__init__(datetime_class)
+
+            self._started: bool = False
+            self._traveller: time_machine.travel | None = None
+            self._coordinates: time_machine.Coordinates | None = None
+
+        def freeze(self) -> Traveller:
+            if self._started:
+                cast(time_machine.Coordinates, self._coordinates).move_to(
+                    self._datetime_class.now(), tick=False
+                )
+            else:
+                self._start(freeze=True)
+
+            return self
+
+        def travel_back(self) -> Traveller:
+            if not self._started:
+                return self
+
+            cast(time_machine.travel, self._traveller).stop()
+            self._coordinates = None
+            self._traveller = None
+            self._started = False
+
+            return self
+
+        def travel(
+            self,
+            years: int = 0,
+            months: int = 0,
+            weeks: int = 0,
+            days: int = 0,
+            hours: int = 0,
+            minutes: int = 0,
+            seconds: int = 0,
+            microseconds: int = 0,
+            *,
+            freeze: bool = False,
+        ) -> Traveller:
+            self._start(freeze=freeze)
+
+            cast(time_machine.Coordinates, self._coordinates).move_to(
+                self._datetime_class.now().add(
+                    years=years,
+                    months=months,
+                    weeks=weeks,
+                    days=days,
+                    hours=hours,
+                    minutes=minutes,
+                    seconds=seconds,
+                    microseconds=microseconds,
+                )
+            )
+
+            return self
+
+        def travel_to(self, dt: DateTime, *, freeze: bool = False) -> Traveller:
+            self._start(freeze=freeze)
+
+            cast(time_machine.Coordinates, self._coordinates).move_to(dt)
+
+            return self
+
+        def _start(self, freeze: bool = False) -> None:
+            if self._started:
+                return
+
+            if not self._traveller:
+                self._traveller = time_machine.travel(
+                    self._datetime_class.now(), tick=not freeze
+                )
+
+            self._coordinates = self._traveller.start()
+
+            self._started = True
+
+        def __enter__(self) -> Traveller:
+            self._start()
+
+            return self
+
+        def __exit__(
+            self,
+            exc_type: type[BaseException] | None,
+            exc_val: BaseException | None,
+            exc_tb: TracebackType,
+        ) -> None:
+            self.travel_back()
+
+else:
+
+    class Traveller(BaseTraveller):  # type: ignore[no-redef]
+
+        ...

pendulum/testing/traveller.py

@@ -28,6 +28,7 @@ include = [
 python = "^3.7"
 python-dateutil = "^2.6"
 "backports.zoneinfo" = { version = "^0.2.1", python = ">=3.7,<3.9" }
+time-machine = { version = "^2.6.0", markers = "implementation_name != 'pypy'" }
 tzdata = ">=2020.1"
 importlib-resources = { version = "^5.9.0", python = ">=3.7,<3.9" }
 

poetry.lock

@@ -6,16 +6,21 @@
 
 import pytest
 import pytz
-import time_machine
 
 from dateutil import tz
 
 import pendulum
 
 from pendulum import DateTime
 from pendulum.tz import timezone
+from pendulum.utils._compat import PYPY
 from tests.conftest import assert_datetime
 
+if not PYPY:
+    import time_machine
+else:
+    time_machine = None
+
 
 @pytest.fixture(autouse=True)
 def _setup():
@@ -105,44 +110,43 @@ def test_now():
     assert now.hour != in_paris.hour
 
 
-@time_machine.travel("2016-03-27 00:30:00Z", tick=False)
-def test_now_dst_off():
-    utc = pendulum.now("UTC")
-    in_paris = pendulum.now("Europe/Paris")
-    in_paris_from_utc = utc.in_tz("Europe/Paris")
-    assert in_paris.hour == 1
-    assert not in_paris.is_dst()
-    assert in_paris.isoformat() == in_paris_from_utc.isoformat()
-
-
-@time_machine.travel("2016-03-27 01:30:00Z", tick=False)
-def test_now_dst_transitioning_on():
-    utc = pendulum.now("UTC")
-    in_paris = pendulum.now("Europe/Paris")
-    in_paris_from_utc = utc.in_tz("Europe/Paris")
-    assert in_paris.hour == 3
-    assert in_paris.is_dst()
-    assert in_paris.isoformat() == in_paris_from_utc.isoformat()
-
-
-@time_machine.travel("2016-10-30 00:30:00Z", tick=False)
-def test_now_dst_on():
-    utc = pendulum.now("UTC")
-    in_paris = pendulum.now("Europe/Paris")
-    in_paris_from_utc = utc.in_tz("Europe/Paris")
-    assert in_paris.hour == 2
-    assert in_paris.is_dst()
-    assert in_paris.isoformat() == in_paris_from_utc.isoformat()
-
-
-@time_machine.travel("2016-10-30 01:30:00Z", tick=False)
-def test_now_dst_transitioning_off():
-    utc = pendulum.now("UTC")
-    in_paris = pendulum.now("Europe/Paris")
-    in_paris_from_utc = utc.in_tz("Europe/Paris")
-    assert in_paris.hour == 2
-    assert not in_paris.is_dst()
-    assert in_paris.isoformat() == in_paris_from_utc.isoformat()
+if time_machine:
+
+    @time_machine.travel("2016-03-27 00:30:00Z", tick=False)
+    def test_now_dst_off():
+        utc = pendulum.now("UTC")
+        in_paris = pendulum.now("Europe/Paris")
+        in_paris_from_utc = utc.in_tz("Europe/Paris")
+        assert in_paris.hour == 1
+        assert not in_paris.is_dst()
+        assert in_paris.isoformat() == in_paris_from_utc.isoformat()
+
+    @time_machine.travel("2016-03-27 01:30:00Z", tick=False)
+    def test_now_dst_transitioning_on():
+        utc = pendulum.now("UTC")
+        in_paris = pendulum.now("Europe/Paris")
+        in_paris_from_utc = utc.in_tz("Europe/Paris")
+        assert in_paris.hour == 3
+        assert in_paris.is_dst()
+        assert in_paris.isoformat() == in_paris_from_utc.isoformat()
+
+    @time_machine.travel("2016-10-30 00:30:00Z", tick=False)
+    def test_now_dst_on():
+        utc = pendulum.now("UTC")
+        in_paris = pendulum.now("Europe/Paris")
+        in_paris_from_utc = utc.in_tz("Europe/Paris")
+        assert in_paris.hour == 2
+        assert in_paris.is_dst()
+        assert in_paris.isoformat() == in_paris_from_utc.isoformat()
+
+    @time_machine.travel("2016-10-30 01:30:00Z", tick=False)
+    def test_now_dst_transitioning_off():
+        utc = pendulum.now("UTC")
+        in_paris = pendulum.now("Europe/Paris")
+        in_paris_from_utc = utc.in_tz("Europe/Paris")
+        assert in_paris.hour == 2
+        assert not in_paris.is_dst()
+        assert in_paris.isoformat() == in_paris_from_utc.isoformat()
 
 
 def test_now_with_fixed_offset():