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

1from dataclasses import dataclass 

2from typing import TYPE_CHECKING, Optional, 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 LargeWtRuleConstraints(BatchingConstraints, JSONWizard): 

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

16 

17 min_wt: Optional[int] 

18 max_wt: Optional[int] 

19 

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

21 key_transform_with_dump = "SNAKE" 

22 tag = RULE_TYPE.LARGE_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.LARGE_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 

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

37 # if ( 

38 # self.min_wt 

39 # and self.min_wt > 0 

40 # and ( 

41 # firing_rule.comparison == COMPARATOR.LESS_THEN 

42 # or firing_rule.comparison == COMPARATOR.LESS_THEN_OR_EQUAL 

43 # ) 

44 # ): 

45 # return False 

46 

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

48 # if self.max_wt and ( 

49 # firing_rule.comparison == COMPARATOR.GREATER_THEN 

50 # or firing_rule.comparison == COMPARATOR.GREATER_THEN_OR_EQUAL 

51 # ): 

52 # return False 

53 

54 return not ( 

55 self.min_wt and firing_rule.value < self.min_wt or self.max_wt and firing_rule.value > self.max_wt 

56 ) 

57 

58 

59def is_large_wt_constraint( 

60 val: BatchingConstraints, 

61) -> TypeGuard[LargeWtRuleConstraints]: 

62 """Check if the constraint is a large waiting time constraint.""" 

63 return val.rule_type == RULE_TYPE.LARGE_WT