Coverage for jutil/command.py: 79%

111 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-10-07 16:40 -0500

1import logging 

2import re 

3import traceback 

4from datetime import datetime, timedelta 

5from typing import Tuple, List, Any, Optional 

6from django.core.management import get_commands, load_command_class 

7from django.core.management.base import BaseCommand, CommandParser 

8from django.utils.timezone import now 

9from django.conf import settings 

10from jutil.dates import ( 

11 last_month, 

12 yesterday, 

13 TIME_RANGE_NAMES, 

14 TIME_STEP_NAMES, 

15 this_month, 

16 last_year, 

17 last_week, 

18 localize_time_range, 

19 this_year, 

20 this_week, 

21 get_time_steps, 

22 next_year, 

23 next_month, 

24 next_week, 

25) 

26from jutil.email import send_email 

27import getpass 

28from django.utils import translation 

29from jutil.parse import parse_datetime 

30 

31 

32logger = logging.getLogger(__name__) 

33 

34 

35class SafeCommand(BaseCommand): 

36 """ 

37 BaseCommand which activates LANGUAGE_CODE locale, catches, logs and emails errors. 

38 Uses list of emails from settings.ADMINS. 

39 Implement do() in derived classes with identical args as normal handle(). 

40 """ 

41 

42 def handle(self, *args, **kwargs): 

43 try: 

44 if hasattr(settings, "LANGUAGE_CODE"): 44 ↛ 46line 44 didn't jump to line 46, because the condition on line 44 was never false

45 translation.activate(settings.LANGUAGE_CODE) 

46 return self.do(*args, **kwargs) 

47 except Exception as e: 

48 msg = "ERROR: {}\nargs: {}\nkwargs: {}\n{}".format(str(e), args, kwargs, traceback.format_exc()) 

49 logger.error(msg) 

50 if not settings.DEBUG: 

51 send_email(settings.ADMINS, "Error @ {}".format(getpass.getuser()), msg) 

52 raise 

53 

54 def do(self, *args, **kwargs): 

55 pass 

56 

57 

58def add_date_range_arguments(parser: CommandParser): 

59 """ 

60 Adds following arguments to the CommandParser: 

61 

62 Ranges: 

63 --begin BEGIN 

64 --end END 

65 --last-year 

66 --last-month 

67 --last-week 

68 --this-year 

69 --this-month 

70 --this-week 

71 --yesterday 

72 --today 

73 --prev-90d 

74 --plus-minus-90d 

75 --next-90d 

76 --prev-60d 

77 --plus-minus-60d 

78 --next-60d 

79 --prev-30d 

80 --plus-minus-30d 

81 --next-30d 

82 --prev-15d 

83 --plus-minus-15d 

84 --next-15d 

85 --prev-7d 

86 --plus-minus-7d 

87 --next-7d 

88 --prev-2d 

89 --plus-minus-2d 

90 --next-2d 

91 --prev-1d 

92 --plus-minus-1d 

93 --next-1d 

94 

95 Steps: 

96 --daily 

97 --weekly 

98 --monthly 

99 

100 :param parser: 

101 :return: 

102 """ 

103 parser.add_argument("--begin", type=str) 

104 parser.add_argument("--end", type=str) 

105 for v in TIME_STEP_NAMES: 

106 parser.add_argument("--" + v.replace("_", "-"), action="store_true") 

107 for v in TIME_RANGE_NAMES: 

108 parser.add_argument("--" + v.replace("_", "-"), action="store_true") 

109 

110 

111def get_date_range_by_name(name: str, today: Optional[datetime] = None, tz: Any = None) -> Tuple[datetime, datetime]: # noqa 

112 """ 

113 Returns a timezone-aware date range by symbolic name. 

114 :param name: Name of the date range. See add_date_range_arguments(). 

115 :param today: Optional current datetime. Default is datetime.utcnow(). 

116 :param tz: Optional timezone. Default is UTC. 

117 :return: datetime (begin, end) 

118 """ 

119 if today is None: 

120 today = datetime.utcnow() 

121 begin = today.replace(hour=0, minute=0, second=0, microsecond=0) 

122 

123 if name == "last_week": 

124 return last_week(today, tz) 

125 if name == "last_month": 

126 return last_month(today, tz) 

127 if name == "last_year": 

128 return last_year(today, tz) 

129 if name == "this_week": 

130 return this_week(today, tz) 

131 if name == "this_month": 

132 return this_month(today, tz) 

133 if name == "this_year": 

134 return this_year(today, tz) 

135 if name == "next_week": 135 ↛ 136line 135 didn't jump to line 136, because the condition on line 135 was never true

136 return next_week(today, tz) 

137 if name == "next_month": 137 ↛ 138line 137 didn't jump to line 138, because the condition on line 137 was never true

138 return next_month(today, tz) 

139 if name == "next_year": 139 ↛ 140line 139 didn't jump to line 140, because the condition on line 139 was never true

140 return next_year(today, tz) 

141 if name == "yesterday": 

142 return yesterday(today, tz) 

143 if name == "today": 

144 return localize_time_range(begin, begin + timedelta(hours=24), tz) 

145 if name == "tomorrow": 145 ↛ 146line 145 didn't jump to line 146, because the condition on line 145 was never true

146 return localize_time_range(begin + timedelta(hours=24), begin + timedelta(hours=48), tz) 

147 

148 m = re.match(r"^plus_minus_(\d+)d$", name) 

149 if m: 

150 days = int(m.group(1)) 

151 return localize_time_range(begin - timedelta(days=days), today + timedelta(days=days), tz) 

152 

153 m = re.match(r"^prev_(\d+)d$", name) 

154 if m: 

155 days = int(m.group(1)) 

156 return localize_time_range(begin - timedelta(days=days), today, tz) 

157 

158 m = re.match(r"^next_(\d+)d$", name) 

159 if m: 159 ↛ 163line 159 didn't jump to line 163, because the condition on line 159 was never false

160 days = int(m.group(1)) 

161 return localize_time_range(begin, today + timedelta(days=days), tz) 

162 

163 raise ValueError("Invalid date range name: {}".format(name)) 

164 

165 

166def parse_date_range_arguments(options: dict, default_range: str = "last_month", tz: Any = None) -> Tuple[datetime, datetime, List[Tuple[datetime, datetime]]]: 

167 """ 

168 Parses date range from input and returns timezone-aware date range and 

169 interval list according to 'step' name argument (optional). 

170 See add_date_range_arguments() 

171 :param options: Parsed arguments passed to the command 

172 :param default_range: Default datetime range to return if no other selected 

173 :param tz: Optional timezone to use. Default is UTC. 

174 :return: begin, end, [(begin1,end1), (begin2,end2), ...] 

175 """ 

176 begin, end = get_date_range_by_name(default_range, tz=tz) 

177 for range_name in TIME_RANGE_NAMES: 

178 if options.get(range_name): 178 ↛ 179line 178 didn't jump to line 179, because the condition on line 178 was never true

179 begin, end = get_date_range_by_name(range_name, tz=tz) 

180 if options.get("begin"): 180 ↛ 183line 180 didn't jump to line 183, because the condition on line 180 was never false

181 begin = parse_datetime(options["begin"], tz) # type: ignore 

182 end = now() 

183 if options.get("end"): 183 ↛ 186line 183 didn't jump to line 186, because the condition on line 183 was never false

184 end = parse_datetime(options["end"], tz) # type: ignore 

185 

186 step_type = "" 

187 for step_name in TIME_STEP_NAMES: 

188 if options.get(step_name): 188 ↛ 189line 188 didn't jump to line 189, because the condition on line 188 was never true

189 if step_type: 

190 raise ValueError("Cannot use --{} and --{} simultaneously".format(step_type, step_name)) 

191 step_type = step_name 

192 if step_type: 192 ↛ 193line 192 didn't jump to line 193, because the condition on line 192 was never true

193 steps = get_time_steps(step_type, begin, end) 

194 else: 

195 steps = [(begin, end)] 

196 return begin, end, steps 

197 

198 

199def get_command_by_name(command_name: str) -> BaseCommand: 

200 """ 

201 Gets Django management BaseCommand derived command class instance by name. 

202 """ 

203 all_commands = get_commands() 

204 app_name = all_commands.get(command_name) 

205 if app_name is None: 205 ↛ 206line 205 didn't jump to line 206, because the condition on line 205 was never true

206 raise Exception(f"Django management command {command_name} not found") 

207 command = app_name if isinstance(app_name, BaseCommand) else load_command_class(app_name, command_name) 

208 assert isinstance(command, BaseCommand) 

209 return command 

210 

211 

212def get_command_name(command: BaseCommand) -> str: 

213 """ 

214 Gets Django management BaseCommand name from instance. 

215 """ 

216 module_name = command.__class__.__module__ 

217 res = module_name.rsplit(".", 1) 

218 if len(res) != 2: 218 ↛ 219line 218 didn't jump to line 219, because the condition on line 218 was never true

219 raise Exception(f"Failed to parse Django command name from {module_name}") 

220 return res[1]