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

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 NotImplementedError() 

22 

23 def travel_back(self) -> Self: 

24 raise NotImplementedError() 

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 NotImplementedError() 

38 

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

40 raise NotImplementedError() 

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 

54if not PYPY: 

55 import time_machine 

56 

57 class Traveller(BaseTraveller): 

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

59 super().__init__(datetime_class) 

60 

61 self._started: bool = False 

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

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

64 

65 def freeze(self) -> Self: 

66 if self._started: 

67 cast(time_machine.Coordinates, self._coordinates).move_to( 

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

69 ) 

70 else: 

71 self._start(freeze=True) 

72 

73 return self 

74 

75 def travel_back(self) -> Self: 

76 if not self._started: 

77 return self 

78 

79 cast(time_machine.travel, self._traveller).stop() 

80 self._coordinates = None 

81 self._traveller = None 

82 self._started = False 

83 

84 return self 

85 

86 def travel( 

87 self, 

88 years: int = 0, 

89 months: int = 0, 

90 weeks: int = 0, 

91 days: int = 0, 

92 hours: int = 0, 

93 minutes: int = 0, 

94 seconds: int = 0, 

95 microseconds: int = 0, 

96 *, 

97 freeze: bool = False, 

98 ) -> Self: 

99 self._start(freeze=freeze) 

100 

101 cast(time_machine.Coordinates, self._coordinates).move_to( 

102 self._datetime_class.now().add( 

103 years=years, 

104 months=months, 

105 weeks=weeks, 

106 days=days, 

107 hours=hours, 

108 minutes=minutes, 

109 seconds=seconds, 

110 microseconds=microseconds, 

111 ) 

112 ) 

113 

114 return self 

115 

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

117 self._start(freeze=freeze) 

118 

119 cast(time_machine.Coordinates, self._coordinates).move_to(dt) 

120 

121 return self 

122 

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

124 if self._started: 

125 return 

126 

127 if not self._traveller: 

128 self._traveller = time_machine.travel( 

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

130 ) 

131 

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

133 

134 self._started = True 

135 

136 def __enter__(self) -> Self: 

137 self._start() 

138 

139 return self 

140 

141 def __exit__( 

142 self, 

143 exc_type: type[BaseException] | None, 

144 exc_val: BaseException | None, 

145 exc_tb: TracebackType, 

146 ) -> None: 

147 self.travel_back() 

148 

149else: 

150 

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

152 ...