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

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 )