Coverage for o2/models/constraints/ready_wt_rule_constraints.py: 95%

1from dataclasses import dataclass 

2from typing import TYPE_CHECKING, TypeGuard 

3 

4from dataclass_wizard import JSONWizard 

5 

6from o2.models.constraints.batching_constraints import BatchingConstraints 

7from o2.models.timetable import RULE_TYPE, FiringRule 

8 

9if TYPE_CHECKING: 

10 from o2.models.timetable import TimetableType 

11 

12 

13@dataclass(frozen=True) 

14class ReadyWtRuleConstraints(BatchingConstraints, JSONWizard): 

15 """Ready waiting time rule constraints for batching.""" 

16 

17 min_wt: int 

18 max_wt: int 

19 

20 class _(JSONWizard.Meta): # noqa: N801 

21 key_transform_with_dump = "SNAKE" 

22 tag = RULE_TYPE.READY_WT.value 

23 tag_key = "rule_type" 

24 

25 def verify_timetable(self, timetable: "TimetableType") -> bool: 

26 """Check if the timetable is valid against the constraints.""" 

27 rules: list[FiringRule[int]] = timetable.get_firing_rules_for_tasks( 

28 self.tasks, rule_type=RULE_TYPE.READY_WT, batch_type=self.batch_type 

29 ) 

30 return all(self._verify_firing_rule(rule) for rule in rules) 

31 

32 def _verify_firing_rule(self, firing_rule: FiringRule[int]) -> bool: 

33 """Check if the firing rule is valid against the constraints.""" 

34 # TODO: We need to do checking in firing_rule pairs 

35 # # Special case: If the rule is a <(=) min wt, it is invalid if min wt is set 

36 # if ( 

37 # self.min_wt 

38 # and self.min_wt > 0 

39 # and ( 

40 # firing_rule.comparison == COMPARATOR.LESS_THEN 

41 # or firing_rule.comparison == COMPARATOR.LESS_THEN_OR_EQUAL 

42 # ) 

43 # ): 

44 # return False 

45 

46 # # Special case: If the rule is a >(=) max wt, it is invalid if max wt is set 

47 # if self.max_wt and ( 

48 # firing_rule.comparison == COMPARATOR.GREATER_THEN 

49 # or firing_rule.comparison == COMPARATOR.GREATER_THEN_OR_EQUAL 

50 # ): 

51 # return False 

52 

53 # if (self.min_wt and firing_rule.value < self.min_wt) or ( 

54 # self.max_wt and firing_rule.value > self.max_wt 

55 # ): 

56 # return False 

57 

58 return True 

59 

60 

61def is_ready_wt_constraint( 

62 val: BatchingConstraints, 

63) -> TypeGuard[ReadyWtRuleConstraints]: 

64 """Check if the constraint is a ready waiting time constraint.""" 

65 return val.rule_type == RULE_TYPE.READY_WT