Coverage for src/dtexp/conditions.py: 100%

1"""Find a date in past/future fulfilling a condition."""

2

3 import calendar

4 import datetime

5 from typing import Literal

6

7 from dtexp.addition import apply_addition_via_pendulum

8 from dtexp.exceptions import DtexpMaxIterationError, DtexpParsingError

9

10 CONDITIONAL_DIRECTION_OPERATORS = {"next", "last", "upcoming", "previous"}

11

12

13 def weekday_occurrence_in_month_from_start(dt: datetime.datetime) -> int:

14 """Return n where dt is the n'th occurrence of its weekday from month start.

15

16 E.g. if its 2025-10-25 then this Saturday is the 4th occurence of this weekday

17 in the month.

18 """

19

20 return ((dt.day - 1) // 7) + 1

21

22

23 def weekday_occurrence_in_month_from_end(dt: datetime.datetime) -> int:

24 """Return n where dt is the n'th occurrence of its weekday from month end.

25

26 E.g. if its 2025-10-25 then this Saturday is the 1st occurence of this weekday

27 from the end of the month.

28 """

29 last_day = calendar.monthrange(dt.year, dt.month)[1]

30

31 # Days remaining in month (including current day)

32 days_from_end = last_day - dt.day

33

34 return ((days_from_end - 1) // 7) + 1

35

36

37 def extract_condition_elements(elements: list[str]) -> tuple[list[str], list[str]]:

38 """Get condition elements until a non-condition-element occurs.

39

40 Returns pair of lists:

41 * First contains the condition elements

42 * second contains the remaining elements

43 """

44 condition_elements = []

45 for element in elements:

46 if element in {

47 "+",

48 "plus",

49 "-",

50 "minus",

51 "/",

52 "next",

53 "past",

54 "upcoming",

55 "previous",

56 }:

57 break

58

59 condition_elements.append(element)

60

61 return condition_elements, elements[len(condition_elements) :]

62

63

64 def build_condition_dict(condition_elements: list[str]) -> dict[str, int]:

65 """Gather conditions fromt the condition elements.

66

67 Returns a dictionary of form {<time_unit>: number}, meaning that

68 this time_unit's value of the datetime must equal the number.

69

70 The entries of this dict later can be logically combined (e.g. with "and"),

71 when evaluating the conditions.

72 """

73 if condition_elements[0] != "where":

74 raise DtexpParsingError("Conditions must start with where")

75

76 cond_dict_equals: dict[str, int] = {}

77 remaining_condition_elements = condition_elements

78

79 while len(remaining_condition_elements) > 0:

80 if remaining_condition_elements[0] in {"where", "and"}:

81 if len(remaining_condition_elements) < 4:

82 msg = (

83 f"Condition expression must consist of 4 elements."

84 f" Got {remaining_condition_elements}"

85 )

86 raise DtexpParsingError(msg)

87

88 condition_unit = remaining_condition_elements[1]

89 if condition_unit not in {

90 "us",

91 "s",

92 "min",

93 "h",

94 "d",

95 "wd",

96 "m",

97 "y",

98 "wdofms",

99 "wdofme",

100 }:

101 msg = f"Cannot understand condition unit {condition_unit}"

102 raise DtexpParsingError(msg)

103

104 condition_type = remaining_condition_elements[2]

105 if condition_type not in {"is"}:

106 raise DtexpParsingError('Can only understand "is" conditions')

107

108 condition_value = int(remaining_condition_elements[3])

109 # Note: We do not check whether the condition_value fits to the condition_type

110 # (i.e. weekdays only 0-6 etc., >0, ...). Instead, this will lead to just not

111 # finding a result

112

113 if cond_dict_equals.get(condition_unit) is None:

114 cond_dict_equals[condition_unit] = condition_value

115 elif cond_dict_equals[condition_unit] != condition_value:

116 # We already have that condition

117 msg = (

118 f"Cannot have two conditions for unit {condition_unit}"

119 " with two different values"

120 f" ({cond_dict_equals[condition_unit]}, {condition_value})"

121 )

122 raise DtexpParsingError(msg)

123

124 remaining_condition_elements = remaining_condition_elements[4:]

125 return cond_dict_equals

126

127

128 def eval_condition(dt: datetime.datetime, condition_elements: list[str]) -> bool:

129 """Evaluate condition on datetime.

130

131 Currently only accepts logical conjunction (and) of conditions.

132 """

133

134 cond_dict_equals = build_condition_dict(condition_elements)

135

136 return (

137 ((conv_val_days := cond_dict_equals.get("d")) is None or dt.day == conv_val_days)

138 and ((conv_val_mins := cond_dict_equals.get("min")) is None or dt.minute == conv_val_mins)

139 and ((conv_val_hours := cond_dict_equals.get("h")) is None or dt.hour == conv_val_hours)

140 and (

141 (conv_val_seconds := cond_dict_equals.get("s")) is None or dt.second == conv_val_seconds

142 )

143 and (

144 (conv_val_weekday := cond_dict_equals.get("wd")) is None

145 or dt.weekday() == conv_val_weekday

146 )

147 and (

148 (conv_val_weekday_occurence_from_month_start := cond_dict_equals.get("wdofms")) is None

149 or weekday_occurrence_in_month_from_start(dt)

150 == conv_val_weekday_occurence_from_month_start

151 )

152 and (

153 (conv_val_weekday_occurence_from_month_end := cond_dict_equals.get("wdofme")) is None

154 or weekday_occurrence_in_month_from_end(dt) == conv_val_weekday_occurence_from_month_end

155 )

156 and ((conv_val_month := cond_dict_equals.get("m")) is None or dt.month == conv_val_month)

157 and ((conv_val_year := cond_dict_equals.get("y")) is None or dt.year == conv_val_year)

158 )

159

160

161 def find_date_by_condition(

162 start_date: datetime.datetime,

163 increment_direction: Literal[1, -1],

164 start_increment: Literal[0, 1],

165 target_unit: str,

166 condition_elements: list[str],

167 max_iter: int = 1000,

168) -> datetime.datetime:

169 """Find a date from a start date in one direction fulfilling a condition.

170

171 Currently this simply iterates in the desired direction and checks the conditions

172 at each step, i.e. brute force search.

173 """

174

175 iter_addition_params = {

176 "microseconds": increment_direction if target_unit == "us" else 0,

177 "seconds": increment_direction if target_unit == "s" else 0,

178 "minutes": increment_direction if target_unit == "min" else 0,

179 "hours": increment_direction if target_unit == "h" else 0,

180 "days": increment_direction if target_unit in {"d", "wd"} else 0,

181 "months": increment_direction if target_unit == "m" else 0,

182 "years": increment_direction if target_unit == "y" else 0,

183 }

184

185 iteration_date = (

186 start_date

187 if start_increment == 0

188 else apply_addition_via_pendulum(start_date, **iter_addition_params)

189 )

190

191 for _ in range(max_iter):

192 if eval_condition(iteration_date, condition_elements):

193 return iteration_date

194

195 iteration_date = apply_addition_via_pendulum(iteration_date, **iter_addition_params)

196

197 msg = f"Max condition iteration limit {max_iter} reached."

198 raise DtexpMaxIterationError(msg)

199

200

201 def handle_condition_expression(

202 split_result: list[str], left: datetime.datetime, max_iter: int = 1000

203) -> tuple[datetime.datetime, list[str]]:

204 """Handle condition expressions.

205

206 Expects split_results to start with a condition expression.

207

208 Will then resolve from left by iterating and checking the conditions up

209 to max_iter steps.

210

211 Returns a pair consisting of the resulting datetime and the list of remaining

212 elements of split_result.

213 """

214 if len(split_result) < 2:

215 raise DtexpParsingError("Condition expression too short")

216

217 operator = split_result[0]

218

219 if operator not in CONDITIONAL_DIRECTION_OPERATORS:

220 msg = f"Condition expression must start with one of {CONDITIONAL_DIRECTION_OPERATORS}"

221 raise DtexpParsingError(msg)

222

223 increment_direction: Literal[-1, 1] = 1 if operator in {"next", "upcoming"} else -1

224 start_increment: Literal[0, 1] = 0 if operator in {"next", "last"} else 1

225

226 target_unit = split_result[1]

227

228 if target_unit not in {"d", "wd", "w", "h", "min", "us", "s", "m", "y"}:

229 msg = f"Unknown target unit {target_unit}."

230 raise DtexpParsingError(msg)

231

232 condition_elements, remaining_elements = extract_condition_elements(split_result[2:])

233

234 result_date = find_date_by_condition(

235 start_date=left,

236 increment_direction=increment_direction,

237 start_increment=start_increment,

238 target_unit=target_unit,

239 condition_elements=condition_elements,

240 max_iter=max_iter,

241 )

242

243 return result_date, remaining_elements