Merge pull request #252 from randovania/bmsbk

BMSBK - Minor cleanup and add functions

src/mercury_engine_data_structures/formats/bmsbk.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import functools
+from enum import Enum
 from typing import TYPE_CHECKING
 
-from construct import (
+import construct
+from construct.core import (
     Array,
     Const,
     Construct,
@@ -13,7 +15,6 @@
     Int32ul,
     Rebuild,
     Struct,
-    Terminated,
 )
 
 from mercury_engine_data_structures.base_resource import BaseResource
@@ -43,7 +44,7 @@ def _rebuild_types(ctx: Container) -> int:
 BlockGroup = Struct(
     "_num_blocks" / Rebuild(Int32ul, _rebuild_blocks),
     "_num_types" / Rebuild(Int32ul, _rebuild_types),
-    "unk_bool" / Flag, # always true?
+    "is_enabled" / Flag, # always true?
     "types" / Array(lambda this: this._num_types, Struct(
         "block_type" / StrId,
         "blocks" / make_vector(Block),
@@ -52,18 +53,48 @@ def _rebuild_types(ctx: Container) -> int:
 
 BMSBK = Struct(
     "magic" / Const(b"MSBK"),
-    "version" / VersionAdapter(),
+    "version" / VersionAdapter("1.10.0"),
     "block_groups" / make_vector(BlockGroup),
     "collision_cameras" / make_vector(Struct(
         "name" / StrId,
         "entries" / make_vector(Int32ul),
     )),
-    Terminated,
+    construct.Terminated,
 )  # fmt: skip
 
 
+class BlockType(Enum):
+    POWER_BEAM = "power_beam"
+    BOMB = "bomb"
+    MISSILE = "missile"
+    SUPER_MISSILE = "super_missile"
+    POWER_BOMB = "power_bomb"
+    BABY = "baby"
+    SCREW_ATTACK = "screw_attack"
+    WEIGHT = "weight"
+
+
 class Bmsbk(BaseResource):
     @classmethod
     @functools.lru_cache
     def construct_class(cls, target_game: Game) -> Construct:
         return BMSBK
+
+    def get_block_group(self, block_group: int) -> Container:
+        """Returns a block group by index"""
+        return self.raw.block_groups[block_group]
+
+    def set_block_type(self, block_group: int, block_type: BlockType) -> None:
+        """Change a given block's block_type into another"""
+        assert len(self.get_block_group(block_group).types) == 1
+        self.get_block_group(block_group).types[0].block_type = block_type
+
+    def get_block(self, block_group: int, block_idx: int) -> Container:
+        """Returns a block from a block group by index"""
+        assert len(self.get_block_group(block_group).types) == 1
+        return self.get_block_group(block_group).types[0].blocks[block_idx]
+
+    def set_respawn_time(self, block_group: int, block_idx: int, respawn_time: float) -> None:
+        """Change the respawn time of a given block"""
+        block = self.get_block(block_group, block_idx)
+        block.respawn_time = respawn_time

tests/formats/test_bmsbk.py

@@ -4,7 +4,7 @@
 from tests.test_lib import parse_build_compare_editor
 
 from mercury_engine_data_structures import samus_returns_data
-from mercury_engine_data_structures.formats.bmsbk import Bmsbk
+from mercury_engine_data_structures.formats.bmsbk import BlockType, Bmsbk
 
 sr_missing = [
     "maps/levels/c10_samus/s908_manicminerbotrun/s908_manicminerbotrun.bmsbk",
@@ -17,3 +17,23 @@
 @pytest.mark.parametrize("bmsbk_path", samus_returns_data.all_files_ending_with(".bmsbk", sr_missing))
 def test_bmsbk(samus_returns_tree, bmsbk_path):
     parse_build_compare_editor(Bmsbk, samus_returns_tree, bmsbk_path)
+
+
+@pytest.fixture()
+def surface_bmsbk(samus_returns_tree) -> Bmsbk:
+    return samus_returns_tree.get_parsed_asset("maps/levels/c10_samus/s000_surface/s000_surface.bmsbk", type_hint=Bmsbk)
+
+
+def test_get_block(surface_bmsbk: Bmsbk):
+    block = surface_bmsbk.get_block(0, 1)
+    assert block.respawn_time == 0.0
+
+
+def test_changing_weakness(surface_bmsbk: Bmsbk):
+    surface_bmsbk.set_block_type(1, BlockType.BOMB)
+    assert surface_bmsbk.get_block_group(1).types[0].block_type == BlockType.BOMB
+
+
+def test_respawn_time(surface_bmsbk: Bmsbk):
+    surface_bmsbk.set_respawn_time(0, 0, 5.0)
+    assert surface_bmsbk.get_block(0, 0).respawn_time == 5.0