Merge pull request #85 from sethmachine/modify-unit-triggers

create triggers for modifying energy/hp; add method to extract CHK file from mpq

src/richchk/io/mpq/starcraft_mpq_io.py

@@ -18,6 +18,24 @@ class StarcraftMpqIo:
     def __init__(self, stormlib_wrapper: StormLibWrapper):
         self._stormlib_wrapper = stormlib_wrapper
 
+    def extract_chk_from_mpq(
+        self,
+        path_to_starcraft_mpq_file: str,
+        outfile: str,
+        overwrite_existing: bool = True,
+    ) -> None:
+        if not os.path.exists(path_to_starcraft_mpq_file):
+            raise FileNotFoundError(path_to_starcraft_mpq_file)
+        open_result = self._stormlib_wrapper.open_archive(
+            path_to_starcraft_mpq_file, StormLibArchiveMode.STORMLIB_READ_ONLY
+        )
+        self._stormlib_wrapper.extract_file(
+            open_result,
+            self._CHK_MPQ_PATH,
+            outfile,
+            overwrite_existing=overwrite_existing,
+        )
+
     def read_chk_from_mpq(self, path_to_starcraft_mpq_file: str) -> RichChk:
         if not os.path.exists(path_to_starcraft_mpq_file):
             raise FileNotFoundError(path_to_starcraft_mpq_file)

src/richchk/model/richchk/trig/actions/modify_unit_energy.py

@@ -0,0 +1,48 @@
+import dataclasses
+from abc import ABC
+
+from ...mrgn.rich_location import RichLocation
+from ...unis.unit_id import UnitId
+from ..player_id import PlayerId
+from ..rich_trigger_action import RichTriggerAction, _RichTriggerActionDefaultsBase
+from ..trigger_action_id import TriggerActionId
+
+
+@dataclasses.dataclass(frozen=True)
+class _ModifyUnitEnergyActionBase(RichTriggerAction, ABC):
+    _group: PlayerId
+    _unit: UnitId
+    _amount: int
+    _percent: int
+    _location: RichLocation
+
+    @classmethod
+    def action_id(cls) -> TriggerActionId:
+        return TriggerActionId.MODIFY_UNIT_ENERGY
+
+    @property
+    def group(self) -> PlayerId:
+        return self._group
+
+    @property
+    def unit(self) -> UnitId:
+        return self._unit
+
+    @property
+    def amount(self) -> int:
+        return self._amount
+
+    @property
+    def percent(self) -> int:
+        return self._percent
+
+    @property
+    def location(self) -> RichLocation:
+        return self._location
+
+
+@dataclasses.dataclass(frozen=True)
+class ModifyUnitEnergyAction(
+    _RichTriggerActionDefaultsBase, _ModifyUnitEnergyActionBase
+):
+    pass

src/richchk/model/richchk/trig/actions/modify_unit_hanger.py

@@ -0,0 +1,48 @@
+import dataclasses
+from abc import ABC
+
+from ...mrgn.rich_location import RichLocation
+from ...unis.unit_id import UnitId
+from ..player_id import PlayerId
+from ..rich_trigger_action import RichTriggerAction, _RichTriggerActionDefaultsBase
+from ..trigger_action_id import TriggerActionId
+
+
+@dataclasses.dataclass(frozen=True)
+class _ModifyUnitHangerActionBase(RichTriggerAction, ABC):
+    _group: PlayerId
+    _unit: UnitId
+    _amount: int
+    _hanger_amount: int
+    _location: RichLocation
+
+    @classmethod
+    def action_id(cls) -> TriggerActionId:
+        return TriggerActionId.MODIFY_UNIT_RESOURCE_AMOUNT
+
+    @property
+    def group(self) -> PlayerId:
+        return self._group
+
+    @property
+    def unit(self) -> UnitId:
+        return self._unit
+
+    @property
+    def amount(self) -> int:
+        return self._amount
+
+    @property
+    def hanger_amount(self) -> int:
+        return self._hanger_amount
+
+    @property
+    def location(self) -> RichLocation:
+        return self._location
+
+
+@dataclasses.dataclass(frozen=True)
+class ModifyUnitHangerAction(
+    _RichTriggerActionDefaultsBase, _ModifyUnitHangerActionBase
+):
+    pass

src/richchk/model/richchk/trig/actions/modify_unit_hitpoints_action.py

@@ -0,0 +1,48 @@
+import dataclasses
+from abc import ABC
+
+from ...mrgn.rich_location import RichLocation
+from ...unis.unit_id import UnitId
+from ..player_id import PlayerId
+from ..rich_trigger_action import RichTriggerAction, _RichTriggerActionDefaultsBase
+from ..trigger_action_id import TriggerActionId
+
+
+@dataclasses.dataclass(frozen=True)
+class _ModifyUnitHitpointsActionBase(RichTriggerAction, ABC):
+    _group: PlayerId
+    _unit: UnitId
+    _amount: int
+    _percent: int
+    _location: RichLocation
+
+    @classmethod
+    def action_id(cls) -> TriggerActionId:
+        return TriggerActionId.MODIFY_UNIT_HIT_POINTS
+
+    @property
+    def group(self) -> PlayerId:
+        return self._group
+
+    @property
+    def unit(self) -> UnitId:
+        return self._unit
+
+    @property
+    def amount(self) -> int:
+        return self._amount
+
+    @property
+    def percent(self) -> int:
+        return self._percent
+
+    @property
+    def location(self) -> RichLocation:
+        return self._location
+
+
+@dataclasses.dataclass(frozen=True)
+class ModifyUnitHitpointsAction(
+    _RichTriggerActionDefaultsBase, _ModifyUnitHitpointsActionBase
+):
+    pass

src/richchk/model/richchk/trig/actions/modify_unit_resources.py

@@ -0,0 +1,48 @@
+import dataclasses
+from abc import ABC
+
+from ...mrgn.rich_location import RichLocation
+from ...unis.unit_id import UnitId
+from ..player_id import PlayerId
+from ..rich_trigger_action import RichTriggerAction, _RichTriggerActionDefaultsBase
+from ..trigger_action_id import TriggerActionId
+
+
+@dataclasses.dataclass(frozen=True)
+class _ModifyUnitResourcesActionBase(RichTriggerAction, ABC):
+    _group: PlayerId
+    _unit: UnitId
+    _amount: int
+    _resource_amount: int
+    _location: RichLocation
+
+    @classmethod
+    def action_id(cls) -> TriggerActionId:
+        return TriggerActionId.MODIFY_UNIT_RESOURCE_AMOUNT
+
+    @property
+    def group(self) -> PlayerId:
+        return self._group
+
+    @property
+    def unit(self) -> UnitId:
+        return self._unit
+
+    @property
+    def amount(self) -> int:
+        return self._amount
+
+    @property
+    def resource_amount(self) -> int:
+        return self._resource_amount
+
+    @property
+    def location(self) -> RichLocation:
+        return self._location
+
+
+@dataclasses.dataclass(frozen=True)
+class ModifyUnitResourcesAction(
+    _RichTriggerActionDefaultsBase, _ModifyUnitResourcesActionBase
+):
+    pass

src/richchk/model/richchk/trig/actions/modify_unit_shields.py

@@ -0,0 +1,48 @@
+import dataclasses
+from abc import ABC
+
+from ...mrgn.rich_location import RichLocation
+from ...unis.unit_id import UnitId
+from ..player_id import PlayerId
+from ..rich_trigger_action import RichTriggerAction, _RichTriggerActionDefaultsBase
+from ..trigger_action_id import TriggerActionId
+
+
+@dataclasses.dataclass(frozen=True)
+class _ModifyUnitShieldsActionBase(RichTriggerAction, ABC):
+    _group: PlayerId
+    _unit: UnitId
+    _amount: int
+    _percent: int
+    _location: RichLocation
+
+    @classmethod
+    def action_id(cls) -> TriggerActionId:
+        return TriggerActionId.MODIFY_UNIT_SHIELD_POINTS
+
+    @property
+    def group(self) -> PlayerId:
+        return self._group
+
+    @property
+    def unit(self) -> UnitId:
+        return self._unit
+
+    @property
+    def amount(self) -> int:
+        return self._amount
+
+    @property
+    def percent(self) -> int:
+        return self._percent
+
+    @property
+    def location(self) -> RichLocation:
+        return self._location
+
+
+@dataclasses.dataclass(frozen=True)
+class ModifyUnitShieldsAction(
+    _RichTriggerActionDefaultsBase, _ModifyUnitShieldsActionBase
+):
+    pass

src/richchk/transcoder/richchk/transcoders/trig/actions/modify_unit_energy_transcoder.py

@@ -0,0 +1,74 @@
+"""Decode modify unit energy action."""
+
+from ......model.chk.trig.decoded_trigger_action import DecodedTriggerAction
+from ......model.richchk.mrgn.rich_location import RichLocation
+from ......model.richchk.richchk_decode_context import RichChkDecodeContext
+from ......model.richchk.richchk_encode_context import RichChkEncodeContext
+from ......model.richchk.trig.actions.modify_unit_energy import ModifyUnitEnergyAction
+from ......model.richchk.trig.player_id import PlayerId
+from ......model.richchk.unis.unit_id import UnitId
+from ......util import logger
+from ...helpers.richchk_enum_transcoder import RichChkEnumTranscoder
+from ..rich_trigger_action_transcoder import RichTriggerActionTranscoder
+from ..rich_trigger_action_transcoder_factory import (
+    _RichTriggerActionRegistrableTranscoder,
+)
+
+
+class RichTriggerModifyUnitEnergyActionTranscoder(
+    RichTriggerActionTranscoder[ModifyUnitEnergyAction, DecodedTriggerAction],
+    _RichTriggerActionRegistrableTranscoder,
+    trigger_action_id=ModifyUnitEnergyAction.action_id(),
+):
+    def __init__(self) -> None:
+        self.log = logger.get_logger(
+            RichTriggerModifyUnitEnergyActionTranscoder.__name__
+        )
+
+    def _decode(
+        self,
+        decoded_action: DecodedTriggerAction,
+        rich_chk_decode_context: RichChkDecodeContext,
+    ) -> ModifyUnitEnergyAction:
+        assert decoded_action.action_id == ModifyUnitEnergyAction.action_id().id
+        assert rich_chk_decode_context.rich_mrgn_lookup is not None
+        maybe_location = rich_chk_decode_context.rich_mrgn_lookup.get_location_by_id(
+            decoded_action.location_id
+        )
+        assert isinstance(maybe_location, RichLocation)
+        return ModifyUnitEnergyAction(
+            _group=RichChkEnumTranscoder.decode_enum(
+                decoded_action.first_group, PlayerId
+            ),
+            _unit=RichChkEnumTranscoder.decode_enum(
+                decoded_action.action_argument_type, UnitId
+            ),
+            _amount=decoded_action.quantifier_or_switch_or_order,
+            _percent=decoded_action.second_group,
+            _location=maybe_location,
+        )
+
+    def _encode(
+        self,
+        rich_action: ModifyUnitEnergyAction,
+        rich_chk_encode_context: RichChkEncodeContext,
+    ) -> DecodedTriggerAction:
+        assert rich_chk_encode_context.rich_mrgn_lookup is not None
+        maybe_location_id = rich_chk_encode_context.rich_mrgn_lookup.get_id_by_location(
+            rich_action.location
+        )
+        assert maybe_location_id is not None
+        return DecodedTriggerAction(
+            _location_id=maybe_location_id,
+            _text_string_id=0,
+            _wav_string_id=0,
+            _time=0,
+            _first_group=RichChkEnumTranscoder.encode_enum(rich_action.group),
+            _second_group=rich_action.percent,
+            _action_argument_type=RichChkEnumTranscoder.encode_enum(rich_action.unit),
+            _action_id=rich_action.action_id().id,
+            _quantifier_or_switch_or_order=rich_action.amount,
+            _flags=0,
+            _padding=0,
+            _mask_flag=0,
+        )