Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pendulum/testing/traveller.py: 32%

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 def _not_implemented(self) -> NotImplementedError: 

53 return NotImplementedError() 

54 

55 

56if not PYPY: 

57 try: 

58 import time_machine 

59 except ImportError: 

60 time_machine = None # type: ignore[assignment] 

61 

62 if time_machine is not None: 

63 

64 class Traveller(BaseTraveller): 

65 def __init__(self, datetime_class: type[DateTime] = DateTime) -> None: 

66 super().__init__(datetime_class) 

67 

68 self._started: bool = False 

69 self._traveller: time_machine.travel | None = None 

70 self._coordinates: time_machine.Coordinates | None = None 

71 

72 def freeze(self) -> Self: 

73 if self._started: 

74 cast("time_machine.Coordinates", self._coordinates).move_to( 

75 self._datetime_class.now(), tick=False 

76 ) 

77 else: 

78 self._start(freeze=True) 

79 

80 return self 

81 

82 def travel_back(self) -> Self: 

83 if not self._started: 

84 return self 

85 

86 cast("time_machine.travel", self._traveller).stop() 

87 self._coordinates = None 

88 self._traveller = None 

89 self._started = False 

90 

91 return self 

92 

93 def travel( 

94 self, 

95 years: int = 0, 

96 months: int = 0, 

97 weeks: int = 0, 

98 days: int = 0, 

99 hours: int = 0, 

100 minutes: int = 0, 

101 seconds: int = 0, 

102 microseconds: int = 0, 

103 *, 

104 freeze: bool = False, 

105 ) -> Self: 

106 self._start(freeze=freeze) 

107 

108 cast("time_machine.Coordinates", self._coordinates).move_to( 

109 self._datetime_class.now().add( 

110 years=years, 

111 months=months, 

112 weeks=weeks, 

113 days=days, 

114 hours=hours, 

115 minutes=minutes, 

116 seconds=seconds, 

117 microseconds=microseconds, 

118 ) 

119 ) 

120 

121 return self 

122 

123 def travel_to(self, dt: DateTime, *, freeze: bool = False) -> Self: 

124 self._start(freeze=freeze) 

125 

126 cast("time_machine.Coordinates", self._coordinates).move_to(dt) 

127 

128 return self 

129 

130 def _start(self, freeze: bool = False) -> None: 

131 if self._started: 

132 return 

133 

134 if not self._traveller: 

135 self._traveller = time_machine.travel( 

136 self._datetime_class.now(), tick=not freeze 

137 ) 

138 

139 self._coordinates = self._traveller.start() 

140 

141 self._started = True 

142 

143 def __enter__(self) -> Self: 

144 self._start() 

145 

146 return self 

147 

148 def __exit__( 

149 self, 

150 exc_type: type[BaseException] | None, 

151 exc_val: BaseException | None, 

152 exc_tb: TracebackType, 

153 ) -> None: 

154 self.travel_back() 

155 

156 else: 

157 

158 class Traveller(BaseTraveller): # type: ignore[no-redef] 

159 def _not_implemented(self) -> NotImplementedError: 

160 return NotImplementedError( 

161 "Time travelling is an optional feature. " 

162 'You can add it by installing Pendulum with the "test" extra.' 

163 ) 

164 

165else: 

166 

167 class Traveller(BaseTraveller): # type: ignore[no-redef] 

168 def _not_implemented(self) -> NotImplementedError: 

169 return NotImplementedError( 

170 "Time travelling is not supported on the PyPy Python implementation." 

171 )