Coverage for o2/models/constraints/daily_hour_rule_constraints.py: 88%

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 BATCH_TYPE, RULE_TYPE, FiringRule 

9 

10if TYPE_CHECKING: 

11 from o2.models.timetable import TimetableType 

12 

13 

14@dataclass(frozen=True) 

15class DailyHourRuleConstraints(BatchingConstraints, JSONWizard): 

16 """Daily hour rule constraints for batching.""" 

17 

18 allowed_hours: dict[DAY, list[int]] 

19 

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

21 key_transform_with_dump = "SNAKE" 

22 tag = RULE_TYPE.DAILY_HOUR.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.DAILY_HOUR, 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: Implement this 

35 return True 

36 

37 def allowed_hours_for_day(self, day: DAY) -> list[int]: 

38 """Get the allowed hours for a specific day.""" 

39 return self.allowed_hours.get(day, []) 

40 

41 @staticmethod 

42 def allow_all( 

43 tasks: list[str], batch_type: BATCH_TYPE = BATCH_TYPE.PARALLEL 

44 ) -> "DailyHourRuleConstraints": 

45 """Create a daily hour rule constraint that allows all hours.""" 

46 return DailyHourRuleConstraints( 

47 id="allow_all", 

48 tasks=tasks, 

49 batch_type=BATCH_TYPE.PARALLEL, 

50 rule_type=RULE_TYPE.DAILY_HOUR, 

51 allowed_hours={day: list(range(24)) for day in DAY}, 

52 ) 

53 

54 

55def is_daily_hour_constraint( 

56 val: BatchingConstraints, 

57) -> TypeGuard[DailyHourRuleConstraints]: 

58 """Check if the constraint is a daily hour constraint.""" 

59 return val.rule_type == RULE_TYPE.DAILY_HOUR