Coverage for o2/models/constraints/week_day_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.days import DAY 

8from o2.models.timetable import RULE_TYPE, FiringRule 

9 

10if TYPE_CHECKING: 

11 from o2.models.timetable import TimetableType 

12 

13 

14@dataclass(frozen=True) 

15class WeekDayRuleConstraints(BatchingConstraints, JSONWizard): 

16 """Week day rule constraints for batching.""" 

17 

18 allowed_days: list[DAY] 

19 

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

21 key_transform_with_dump = "SNAKE" 

22 tag = RULE_TYPE.WEEK_DAY.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[DAY]] = timetable.get_firing_rules_for_tasks( 

28 self.tasks, rule_type=RULE_TYPE.WEEK_DAY, 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[DAY]) -> bool: 

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

34 return firing_rule.value in self.allowed_days 

35 

36 

37def is_week_day_constraint( 

38 val: BatchingConstraints, 

39) -> TypeGuard[WeekDayRuleConstraints]: 

40 """Check if the constraint is a week day constraint.""" 

41 return val.rule_type == RULE_TYPE.WEEK_DAY