fix-3

Author: randriu

pyproject.toml

@@ -71,7 +71,7 @@ exclude = [
   ".git",
   "__pycache__",
 ]
-target-version = "py311"
+target-version = "py312"
 
 [tool.ruff.lint]
 select = ["I"]  # enable isort (import sorting) rules

umbi/ats/actions_mixin.py

@@ -0,0 +1,183 @@
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+from .choices_mixin import HasBranchSpace, HasChoiceSpace
+from .entity_space import EntityMapping, OptionalMappingManager
+from .entity_space_mixins import HasBranchActionSpace, HasChoiceActionSpace
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ChoiceToChoiceActionMixin(HasChoiceSpace, HasChoiceActionSpace):
+    #: Choice-to-choice-action mapping manager.
+    _choice_to_choice_action_manager: OptionalMappingManager[int] = field(init=False)
+
+    def __post_init__(self):
+        self._choice_to_choice_action_manager = OptionalMappingManager[int](
+            name="choice_to_choice_action",
+            domain=self._choice_space,
+            can_have_mapping=lambda: self._can_have_choice_to_choice_action(),
+            must_have_mapping=lambda: self._must_have_choice_to_choice_action(),
+        )
+        self._choice_action_space._subscribe(self, lambda *args: self._choice_to_choice_action_manager.auto_manage())
+        super().__post_init__()
+
+    def _can_have_choice_to_choice_action(self) -> bool:
+        return self._choice_action_space.has_entities
+
+    def _must_have_choice_to_choice_action(self) -> bool:
+        return self._can_have_choice_to_choice_action()
+
+    @property
+    def has_choice_to_choice_action(self) -> bool:
+        return self._choice_to_choice_action_manager.has_mapping
+
+    @property
+    def choice_to_choice_action(self) -> EntityMapping[int]:
+        return self._choice_to_choice_action_manager.mapping
+
+    @choice_to_choice_action.setter
+    def choice_to_choice_action(self, values: Sequence[int]) -> None:
+        self._choice_to_choice_action_manager.mapping = values
+
+    def validate(self):
+        self._choice_to_choice_action_manager.validate()
+        super().validate()
+
+
+@dataclass
+class ChoiceActionToNameMixin(HasChoiceActionSpace):
+    #: Choice-to-choice-action mapping manager.
+    _choice_action_to_name_manager: OptionalMappingManager[str] = field(init=False)
+
+    def __post_init__(self):
+        self._choice_action_to_name_manager = OptionalMappingManager[str](
+            name="choice_action_to_name",
+            domain=self._choice_action_space,
+            can_have_mapping=lambda: self._can_have_choice_action_to_name(),
+            must_have_mapping=lambda: self._must_have_choice_action_to_name(),
+        )
+        self._choice_action_space._subscribe(self, lambda *args: self._choice_action_to_name_manager.auto_manage())
+        super().__post_init__()
+
+    def _can_have_choice_action_to_name(self) -> bool:
+        return self._choice_action_space.has_entities
+
+    def _must_have_choice_action_to_name(self) -> bool:
+        return False  # choice_action_to_name is always optional
+
+    def new_choice_action_to_name(self):
+        self._choice_action_to_name_manager.create_mapping()
+
+    def remove_choice_action_to_name(self):
+        self._choice_action_to_name_manager.remove_mapping()
+
+    @property
+    def has_choice_action_to_name(self) -> bool:
+        return self._choice_action_to_name_manager.has_mapping
+
+    @property
+    def choice_action_to_name(self) -> EntityMapping[str]:
+        return self._choice_action_to_name_manager.mapping
+
+    @choice_action_to_name.setter
+    def choice_action_to_name(self, values: Sequence[str] | None) -> None:
+        self._choice_action_to_name_manager.mapping = values
+
+    def validate(self):
+        self._choice_action_to_name_manager.validate()
+        super().validate()
+
+
+@dataclass
+class BranchToBranchActionMixin(HasBranchSpace, HasBranchActionSpace):
+    #: Branch-to-branch-action mapping manager.
+    _branch_to_branch_action_manager: OptionalMappingManager[int] = field(init=False)
+
+    def __post_init__(self):
+        self._branch_to_branch_action_manager = OptionalMappingManager[int](
+            name="branch_to_branch_action",
+            domain=self._branch_space,
+            can_have_mapping=lambda *args: self._can_have_branch_to_branch_action(),
+            must_have_mapping=lambda *args: self._must_have_branch_to_branch_action(),
+        )
+        self._branch_action_space._subscribe(self, lambda *args: self._branch_to_branch_action_manager.auto_manage())
+        super().__post_init__()
+
+    def _can_have_branch_to_branch_action(self) -> bool:
+        return self._branch_action_space.has_entities
+
+    def _must_have_branch_to_branch_action(self) -> bool:
+        return self._can_have_branch_to_branch_action()
+
+    @property
+    def has_branch_to_branch_action(self) -> bool:
+        return self._branch_to_branch_action_manager.has_mapping
+
+    @property
+    def branch_to_branch_action(self) -> EntityMapping[int]:
+        return self._branch_to_branch_action_manager.mapping
+
+    @branch_to_branch_action.setter
+    def branch_to_branch_action(self, values: Sequence[int]) -> None:
+        self._branch_to_branch_action_manager.mapping = values
+
+    def validate(self):
+        self._branch_to_branch_action_manager.validate()
+        super().validate()
+
+
+@dataclass
+class BranchActionToNameMixin(HasBranchActionSpace):
+    #: Branch-to-branch-action mapping manager.
+    _branch_action_to_name_manager: OptionalMappingManager[str] = field(init=False)
+
+    def __post_init__(self):
+        self._branch_action_to_name_manager = OptionalMappingManager[str](
+            name="branch_action_to_name",
+            domain=self._branch_action_space,
+            can_have_mapping=lambda *args: self._can_have_branch_action_to_name(),
+            must_have_mapping=lambda *args: self._must_have_branch_action_to_name(),
+        )
+        self._branch_action_space._subscribe(self, lambda *args: self._branch_action_to_name_manager.auto_manage())
+        super().__post_init__()
+
+    def _can_have_branch_action_to_name(self) -> bool:
+        return self._branch_action_space.has_entities
+
+    def _must_have_branch_action_to_name(self) -> bool:
+        return False  # branch_action_to_name is always optional
+
+    def add_branch_action_to_name(self):
+        self._branch_action_to_name_manager.create_mapping()
+
+    def remove_branch_action_to_name(self):
+        self._branch_action_to_name_manager.remove_mapping()
+
+    @property
+    def has_branch_action_to_name(self) -> bool:
+        return self._branch_action_to_name_manager.has_mapping
+
+    @property
+    def branch_action_to_name(self) -> EntityMapping[str]:
+        return self._branch_action_to_name_manager.mapping
+
+    @branch_action_to_name.setter
+    def branch_action_to_name(self, values: Sequence[str] | None) -> None:
+        self._branch_action_to_name_manager.mapping = values
+
+    def validate(self):
+        self._branch_action_to_name_manager.validate()
+        super().validate()
+
+
+class ActionsMixin(
+    ChoiceToChoiceActionMixin,
+    ChoiceActionToNameMixin,
+    BranchToBranchActionMixin,
+    BranchActionToNameMixin,
+):
+    def __post_init__(self):
+        super().__post_init__()

umbi/ats/annotations_mixin.py

@@ -0,0 +1,117 @@
+import logging
+from dataclasses import dataclass, field
+
+logger = logging.getLogger(__name__)
+
+from .annotations import (
+    Annotation,
+    AtomicPropositionAnnotation,
+    RewardAnnotation,
+)
+from .entity_space_mixins import HasCommonEntitySpaces
+
+
+@dataclass
+class AnnotationsMixin(HasCommonEntitySpaces):
+    #: Annotation category -> (annotation name -> annotation) mapping. Categories 'rewards' and 'aps' can be used, but
+    #: must be of the type RewardAnnotation and AtomicPropositionAnnotation, respectively.
+    _annotations: dict[str, dict[str, Annotation]] = field(default_factory=lambda: {})
+
+    @property
+    def annotations(self) -> dict[str, dict[str, Annotation]]:
+        """Get the annotations of the ATS."""
+        return self._annotations
+
+    @property
+    def annotation_categories(self) -> list[str]:
+        """Get the annotation categories of the ATS."""
+        return list(self.annotations.keys())
+
+    def new_annotation(self, category: str, name: str, **kwargs) -> Annotation:
+        """Add an annotation to the ATS."""
+        if category not in self.annotations:
+            self.annotations[category] = {}
+        if name in self.annotations[category]:
+            raise ValueError(f"annotation with name {name} already exists in category {category}")
+        annotation = Annotation(name=name, entity_spaces=self._common_entity_spaces, **kwargs)
+        self.annotations[category][name] = annotation
+        return annotation
+
+    @property
+    def reward_annotations(self) -> dict[str, RewardAnnotation]:
+        """Get the reward annotations."""
+        if "rewards" not in self.annotations:
+            self.annotations["rewards"] = {}
+        return self.annotations["rewards"]  # type: ignore
+
+    @property
+    def has_reward_annotations(self) -> bool:
+        """Check if there are any reward annotations."""
+        return len(self.reward_annotations) > 0
+
+    @property
+    def reward_annotation_names(self) -> list[str]:
+        """Get the names of all reward annotations."""
+        return list(self.reward_annotations.keys())
+
+    def has_reward_annotation(self, name: str) -> bool:
+        """Check if a reward annotation with the given name exists."""
+        return name in self.reward_annotations
+
+    def new_reward_annotation(self, name: str, **kwargs) -> RewardAnnotation:
+        """Add a reward annotation."""
+        if self.has_reward_annotation(name):
+            raise ValueError(f"reward annotation with name {name} already exists")
+        annotation = RewardAnnotation(name=name, entity_spaces=self._common_entity_spaces, **kwargs)
+        self.reward_annotations[name] = annotation
+        return annotation
+
+    def get_reward_annotation(self, name: str) -> RewardAnnotation:
+        """Get the reward annotation with the given name."""
+        if not self.has_reward_annotation(name):
+            raise ValueError(f"reward annotation with name {name} does not exist")
+        return self.reward_annotations[name]
+
+    ### Atomic propositions. ###
+
+    @property
+    def has_ap_annotations(self) -> bool:
+        """Check if there are any atomic proposition annotations."""
+        return len(self.ap_annotations) > 0
+
+    @property
+    def ap_annotations(self) -> dict[str, AtomicPropositionAnnotation]:
+        """Get the atomic proposition annotations."""
+        if "aps" not in self.annotations:
+            self.annotations["aps"] = {}
+        return self.annotations["aps"]  # type: ignore
+
+    @property
+    def ap_annotation_names(self) -> list[str]:
+        """Get the names of all atomic proposition annotations."""
+        return list(self.ap_annotations.keys())
+
+    def has_ap_annotation(self, name: str) -> bool:
+        """Check if an atomic proposition annotation with the given name exists."""
+        return name in self.ap_annotations
+
+    def new_ap_annotation(self, name: str, **kwargs) -> AtomicPropositionAnnotation:
+        """Add an atomic proposition annotation."""
+        if self.has_ap_annotation(name):
+            raise ValueError(f"atomic proposition annotation with name {name} already exists")
+        annotation = AtomicPropositionAnnotation(name=name, entity_spaces=self._common_entity_spaces, **kwargs)
+        self.ap_annotations[name] = annotation
+        return annotation
+
+    def get_ap_annotation(self, name: str) -> AtomicPropositionAnnotation:
+        """Get the atomic proposition annotation with the given name."""
+        if not self.has_ap_annotation(name):
+            raise ValueError(f"atomic proposition annotation with name {name} does not exist")
+        return self.ap_annotations[name]
+
+    def validate(self):
+        """Validate annotations of the ATS."""
+        for _, annotations in self.annotations.items():
+            for _, annotation in annotations.items():
+                annotation.validate()
+        super().validate()