Coverage for o2/models/rule_selector.py: 80%

1from dataclasses import dataclass 

2from typing import TYPE_CHECKING, Optional 

3 

4from dataclass_wizard import JSONWizard 

5 

6if TYPE_CHECKING: 

7 from o2.models.state import State 

8 from o2.models.timetable import BatchingRule, FiringRule 

9 

10 

11@dataclass(frozen=True) 

12class RuleSelector(JSONWizard): 

13 """Rule selector for batching rule and firing rule.""" 

14 

15 batching_rule_task_id: str 

16 # OR Index, AND Index 

17 firing_rule_index: Optional[tuple[int, int]] 

18 

19 @staticmethod 

20 def from_batching_rule( 

21 batching_rule: "BatchingRule", 

22 firing_rule_index: Optional[tuple[int, int]] = None, 

23 ) -> "RuleSelector": 

24 """Create a rule selector from a batching rule.""" 

25 return RuleSelector( 

26 batching_rule_task_id=batching_rule.task_id, 

27 firing_rule_index=firing_rule_index, 

28 ) 

29 

30 @property 

31 def has_firing_rule(self) -> bool: 

32 """Check if the selector has a firing rule.""" 

33 return self.firing_rule_index is not None 

34 

35 def get_batching_rule_from_state(self, state: "State") -> Optional["BatchingRule"]: 

36 """Get a batching rule by rule selector.""" 

37 return next( 

38 (rule for rule in state.timetable.batch_processing if rule.task_id == self.batching_rule_task_id), 

39 None, 

40 ) 

41 

42 def get_firing_rule_from_state(self, state: "State") -> Optional["FiringRule"]: 

43 """Get a firing rule by rule selector.""" 

44 if self.firing_rule_index is None: 

45 return None 

46 batching_rule = self.get_batching_rule_from_state(state) 

47 if batching_rule is None: 

48 return None 

49 if self.firing_rule_index[0] >= len(batching_rule.firing_rules): 

50 return None 

51 if self.firing_rule_index[1] >= len(batching_rule.firing_rules[self.firing_rule_index[0]]): 

52 return None 

53 return batching_rule.firing_rules[self.firing_rule_index[0]][self.firing_rule_index[1]] 

54 

55 def id(self) -> str: 

56 """Generate a unique identifier for this rule selector. 

57 

58 Creates a string representation that can be used to identify this rule selector. 

59 """ 

60 if self.firing_rule_index is None: 

61 return f"#{self.batching_rule_task_id}" 

62 return f"#{self.batching_rule_task_id}-{self.firing_rule_index[0]}-{self.firing_rule_index[1]}" 

63 

64 def __str__(self) -> str: 

65 """Return a human-readable string representation of the rule selector. 

66 

67 Formats the rule information in a readable format showing the batching rule 

68 and firing rule indices if available. 

69 """ 

70 if self.firing_rule_index is not None: 

71 return ( 

72 f"Batching Rule: {self.batching_rule_task_id}" 

73 f">{self.firing_rule_index[0]}>{self.firing_rule_index[1]}" 

74 ) 

75 return f"Batching Rule: {self.batching_rule_task_id}"