Coverage for o2/actions/batching_actions/remove_date_time_rule_action.py: 97%

1from dataclasses import replace 

2 

3from typing_extensions import Required, override 

4 

5from o2.actions.base_actions.base_action import BaseAction, BaseActionParamsType, RateSelfReturnType 

6from o2.models.days import DAY 

7from o2.models.solution import Solution 

8from o2.models.state import State 

9from o2.store import Store 

10 

11 

12class RemoveDateTimeRuleActionParamsType(BaseActionParamsType): 

13 """Parameter for RemoveDateTimeRuleAction.""" 

14 

15 task_id: Required[str] 

16 day: Required[DAY] 

17 

18 

19class RemoveDateTimeRuleAction(BaseAction): 

20 """RemoveDateTimeRuleAction will remove a day of week and time of day rule.""" 

21 

22 params: RemoveDateTimeRuleActionParamsType 

23 

24 @override 

25 def apply(self, state: State, enable_prints: bool = True) -> State: 

26 timetable = state.timetable 

27 task_id = self.params["task_id"] 

28 

29 best_selector = timetable.get_longest_time_period_for_daily_hour_firing_rules( 

30 task_id, self.params["day"] 

31 ) 

32 

33 if best_selector is None: 

34 return state 

35 

36 day_selector, lower_bound_selector, upper_bound_selector = best_selector 

37 

38 index, batching_rule = timetable.get_batching_rule(lower_bound_selector) 

39 assert batching_rule is not None and index is not None 

40 

41 day_selector_index = ( 

42 day_selector.firing_rule_index[1] 

43 if day_selector is not None and day_selector.firing_rule_index is not None 

44 else -1 

45 ) 

46 lower_bound_index = ( 

47 lower_bound_selector.firing_rule_index[1] 

48 if lower_bound_selector.firing_rule_index is not None 

49 else -1 

50 ) 

51 upper_bound_index = ( 

52 upper_bound_selector.firing_rule_index[1] 

53 if upper_bound_selector.firing_rule_index is not None 

54 else -1 

55 ) 

56 assert upper_bound_selector.firing_rule_index is not None 

57 or_rules_index = upper_bound_selector.firing_rule_index[0] 

58 

59 if or_rules_index >= len(batching_rule.firing_rules): 

60 return state 

61 

62 new_and_rule = [ 

63 and_rule 

64 for i, and_rule in enumerate(batching_rule.firing_rules[or_rules_index]) 

65 if i 

66 not in [ 

67 day_selector_index, 

68 lower_bound_index, 

69 upper_bound_index, 

70 ] 

71 ] 

72 new_batching_rule = replace( 

73 batching_rule, 

74 firing_rules=batching_rule.firing_rules[:or_rules_index] 

75 + ([new_and_rule] if new_and_rule else []) 

76 + batching_rule.firing_rules[or_rules_index + 1 :], 

77 ) 

78 if len(new_batching_rule.firing_rules) == 0: 

79 return state.replace_timetable( 

80 batch_processing=timetable.batch_processing[:index] + timetable.batch_processing[index + 1 :], 

81 ) 

82 

83 return state.replace_timetable( 

84 batch_processing=timetable.batch_processing[:index] 

85 + [new_batching_rule] 

86 + timetable.batch_processing[index + 1 :], 

87 ) 

88 

89 @override 

90 @staticmethod 

91 def rate_self(store: Store, input: Solution) -> RateSelfReturnType: 

92 raise NotImplementedError("Not implemented")