1from __future__ import annotations
2
3from typing import TYPE_CHECKING
4from typing import cast
5
6from pendulum.datetime import DateTime
7from pendulum.utils._compat import PYPY
8
9
10if TYPE_CHECKING:
11 from types import TracebackType
12
13 from typing_extensions import Self
14
15
16class BaseTraveller:
17 def __init__(self, datetime_class: type[DateTime] = DateTime) -> None:
18 self._datetime_class: type[DateTime] = datetime_class
19
20 def freeze(self) -> Self:
21 raise self._not_implemented()
22
23 def travel_back(self) -> Self:
24 raise self._not_implemented()
25
26 def travel(
27 self,
28 years: int = 0,
29 months: int = 0,
30 weeks: int = 0,
31 days: int = 0,
32 hours: int = 0,
33 minutes: int = 0,
34 seconds: int = 0,
35 microseconds: int = 0,
36 ) -> Self:
37 raise self._not_implemented()
38
39 def travel_to(self, dt: DateTime, *, freeze: bool = False) -> Self:
40 raise self._not_implemented()
41
42 def __enter__(self) -> Self:
43 return self
44
45 def __exit__(
46 self,
47 exc_type: type[BaseException] | None,
48 exc_val: BaseException | None,
49 exc_tb: TracebackType,
50 ) -> None:
51 ...
52
53 def _not_implemented(self) -> NotImplementedError:
54 return NotImplementedError()
55
56
57if not PYPY:
58 try:
59 import time_machine
60 except ImportError:
61 time_machine = None # type: ignore[assignment]
62
63 if time_machine is not None:
64
65 class Traveller(BaseTraveller):
66 def __init__(self, datetime_class: type[DateTime] = DateTime) -> None:
67 super().__init__(datetime_class)
68
69 self._started: bool = False
70 self._traveller: time_machine.travel | None = None
71 self._coordinates: time_machine.Coordinates | None = None
72
73 def freeze(self) -> Self:
74 if self._started:
75 cast(time_machine.Coordinates, self._coordinates).move_to(
76 self._datetime_class.now(), tick=False
77 )
78 else:
79 self._start(freeze=True)
80
81 return self
82
83 def travel_back(self) -> Self:
84 if not self._started:
85 return self
86
87 cast(time_machine.travel, self._traveller).stop()
88 self._coordinates = None
89 self._traveller = None
90 self._started = False
91
92 return self
93
94 def travel(
95 self,
96 years: int = 0,
97 months: int = 0,
98 weeks: int = 0,
99 days: int = 0,
100 hours: int = 0,
101 minutes: int = 0,
102 seconds: int = 0,
103 microseconds: int = 0,
104 *,
105 freeze: bool = False,
106 ) -> Self:
107 self._start(freeze=freeze)
108
109 cast(time_machine.Coordinates, self._coordinates).move_to(
110 self._datetime_class.now().add(
111 years=years,
112 months=months,
113 weeks=weeks,
114 days=days,
115 hours=hours,
116 minutes=minutes,
117 seconds=seconds,
118 microseconds=microseconds,
119 )
120 )
121
122 return self
123
124 def travel_to(self, dt: DateTime, *, freeze: bool = False) -> Self:
125 self._start(freeze=freeze)
126
127 cast(time_machine.Coordinates, self._coordinates).move_to(dt)
128
129 return self
130
131 def _start(self, freeze: bool = False) -> None:
132 if self._started:
133 return
134
135 if not self._traveller:
136 self._traveller = time_machine.travel(
137 self._datetime_class.now(), tick=not freeze
138 )
139
140 self._coordinates = self._traveller.start()
141
142 self._started = True
143
144 def __enter__(self) -> Self:
145 self._start()
146
147 return self
148
149 def __exit__(
150 self,
151 exc_type: type[BaseException] | None,
152 exc_val: BaseException | None,
153 exc_tb: TracebackType,
154 ) -> None:
155 self.travel_back()
156
157 else:
158
159 class Traveller(BaseTraveller): # type: ignore[no-redef]
160 def _not_implemented(self) -> NotImplementedError:
161 return NotImplementedError(
162 "Time travelling is an optional feature. "
163 'You can add it by installing Pendulum with the "test" extra.'
164 )
165
166else:
167
168 class Traveller(BaseTraveller): # type: ignore[no-redef]
169 def _not_implemented(self) -> NotImplementedError:
170 return NotImplementedError(
171 "Time travelling is not supported on the PyPy Python implementation."
172 )