Coverage for o2/models/constraints/daily_hour_rule_constraints.py: 88%
25 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-16 11:18 +0000
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-16 11:18 +0000
1from dataclasses import dataclass
2from typing import TYPE_CHECKING, TypeGuard
4from dataclass_wizard import JSONWizard
6from o2.models.constraints.batching_constraints import BatchingConstraints
7from o2.models.days import DAY
8from o2.models.timetable import BATCH_TYPE, RULE_TYPE, FiringRule
10if TYPE_CHECKING:
11 from o2.models.timetable import TimetableType
14@dataclass(frozen=True)
15class DailyHourRuleConstraints(BatchingConstraints, JSONWizard):
16 """Daily hour rule constraints for batching."""
18 allowed_hours: dict[DAY, list[int]]
20 class _(JSONWizard.Meta): # noqa: N801
21 key_transform_with_dump = "SNAKE"
22 tag = RULE_TYPE.DAILY_HOUR.value
23 tag_key = "rule_type"
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)
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
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, [])
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 )
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