Coverage for o2/actions/batching_actions/remove_rule_action.py: 84%

1from dataclasses import dataclass 

2 

3from typing_extensions import override 

4 

5from o2.actions.base_actions.base_action import ( 

6 RateSelfReturnType, 

7) 

8from o2.actions.base_actions.batching_rule_base_action import ( 

9 BatchingRuleBaseAction, 

10 BatchingRuleBaseActionParamsType, 

11) 

12from o2.models.self_rating import RATING 

13from o2.models.settings import Settings 

14from o2.models.solution import Solution 

15from o2.models.state import State 

16from o2.store import Store 

17from o2.util.logger import warn 

18 

19 

20class RemoveRuleActionParamsType(BatchingRuleBaseActionParamsType): 

21 """Parameter for `RemoveRuleAction`.""" 

22 

23 pass 

24 

25 

26@dataclass(frozen=True) 

27class RemoveRuleAction(BatchingRuleBaseAction, str=False): 

28 """`RemoveRuleAction` will remove a `FiringRule` from a `BatchingRule`. 

29 

30 It will not be smart about this, but just yield random firing rules to remove 

31 """ 

32 

33 params: RemoveRuleActionParamsType 

34 

35 @override 

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

37 timetable = state.timetable 

38 rule_selector = self.params["rule"] 

39 

40 index, rule = timetable.get_batching_rule(rule_selector) 

41 if rule is None or index is None: 

42 warn(f"BatchingRule not found for {rule_selector}") 

43 return state 

44 if enable_prints: 

45 warn(f"\t\t>> Removing rule {rule_selector}") 

46 

47 new_batching_rule = rule.remove_firing_rule(rule_selector) 

48 if new_batching_rule is None: 

49 return state.replace_timetable( 

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

51 ) 

52 

53 return state.replace_timetable( 

54 batch_processing=timetable.batch_processing[:index] 

55 + [new_batching_rule] 

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

57 ) 

58 

59 @override 

60 @staticmethod 

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

62 if Settings.DISABLE_REMOVE_ACTION_RULE: 

63 return ( 

64 RATING.NOT_APPLICABLE, 

65 None, 

66 ) 

67 

68 selectors = [ 

69 rule_selector 

70 for batching_rule in store.current_timetable.batch_processing 

71 for rule_selector in batching_rule.get_firing_rule_selectors() 

72 ] 

73 

74 for rule_selector in selectors: 

75 duration_fn = store.constraints.get_duration_fn_for_task(rule_selector.batching_rule_task_id) 

76 yield ( 

77 RATING.VERY_LOW, 

78 RemoveRuleAction(RemoveRuleActionParamsType(rule=rule_selector, duration_fn=duration_fn)), 

79 )