Add more methods to BRSA

src/mercury_engine_data_structures/formats/brsa.py

@@ -3,13 +3,15 @@
 import functools
 from typing import TYPE_CHECKING
 
+from construct import Container, ListContainer
+
 from mercury_engine_data_structures.base_resource import BaseResource
 from mercury_engine_data_structures.formats import standard_format
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
-    from construct import Construct, Container
+    from construct import Construct
 
     from mercury_engine_data_structures.game_check import Game
 
@@ -24,8 +26,168 @@
     def subarea_setups(self) -> Iterator[Container]:
         yield from self.raw.Root.pSubareaManager.vSubareaSetups
 
-    def get_subarea_setup(self, id: str) -> Container:
-        return next(setup for setup in self.subarea_setups if setup.sId == id)
+    @property
+    def charclass_groups(self) -> Iterator[Container]:
+        yield from self.raw.Root.pSubareaManager.vCharclassGroups
+
+    def get_subarea_setup(self, setup_id: str) -> Container:
+        """Gets a setup
+
+        param setup_id: the name of the setup
+        returns: the setup"""
+        return next(setup for setup in self.subarea_setups if setup.sId == setup_id)
+
+    def add_setup(self, setup_id: str) -> Container:
+        """Adds a new setup
+
+        param setup_id: the name of the new setup
+        returns: the newly created setup"""
+        if next((setup for setup in self.subarea_setups if setup.sId == setup_id), None):
+            raise ValueError(f"Setup {setup_id} is already present")
+
+        new_setup = Container({"sId": setup_id, "vSubareaConfigs": ListContainer()})
+        self.raw.Root.pSubareaManager.vSubareaSetups.append(new_setup)
+
+        return new_setup
+
+    def get_subarea_config(self, subarea_id: str, setup_id: str = "Default") -> Container:
+        """Gets a config for a subarea
+
+        param subarea_id: the name of the subarea the config is for
+        param setup_id: the name of the setup the config is in
+        returns: the config for the subarea"""
+        return next(config for config in self.get_subarea_setup(setup_id).vSubareaConfigs if config.sId == subarea_id)
+
+    def add_subarea_config(
+        self,
+        subarea_id: str,
+        setup_id: str = "Default",
+        *,
+        disable_subarea: bool = False,
+        camera_distance: float = -1.0,
+        ignore_camera_offsets: bool = False,
+        charclass_group: str = "No Enemies",
+        camera_ids: list[str] = [],
+        cutscene_ids: list[str] = [],
+    ) -> Container:
+        """Adds a config for a subarea
+
+        param subarea_id: the name of the subarea the config is for
+        param setup_id: the name of the setup the config will be in
+        returns: the newly created config"""
+        if next(
+            (config for config in self.get_subarea_setup(setup_id).vSubareaConfigs if config.sId == subarea_id), None
+        ):
+            raise ValueError(f"Config for {subarea_id} is already present in {setup_id}")
+
+        new_config = Container(
+            {
+                "sId": subarea_id,
+                "sSetupId": setup_id,
+                "bDisableSubarea": disable_subarea,
+                "fCameraZDistance": camera_distance,
+                "bIgnoreMetroidCameraOffsets": ignore_camera_offsets,
+                "sCharclassGroupId": charclass_group,
+                "asItemsIds": ListContainer([""] * 9),
+                "vsCameraCollisionsIds": ListContainer(camera_ids) if camera_ids else ListContainer(),
+                "vsCutscenesIds": ListContainer(cutscene_ids) if cutscene_ids else ListContainer(),
+            }
+        )
+        self.get_subarea_setup(setup_id).vSubareaConfigs.append(new_config)
+
+        return new_config
+
+    def set_scenario_collider(self, subarea_id: str, collider_name: str, setup_id: str = "Default") -> None:
+        """Sets the scenario collider for a subarea
+
+        param subarea_id: the name of the subarea
+        param collider_name: the name of the collider
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[0] = collider_name
+
+    def set_light_group(self, subarea_id: str, group_name: str, setup_id: str = "Default") -> None:
+        """Sets the light group for a subarea
+
+        param subarea_id: the name of the subarea
+        param group_name: the name of the group
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[1] = group_name
+
+    def set_sound_group(self, subarea_id: str, group_name: str, setup_id: str = "Default") -> None:
+        """Sets the sound group for a subarea
+
+        param subarea_id: the name of the subarea
+        param collider_name: the name of the group
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[2] = group_name
+
+    def set_scene_group(self, subarea_id: str, group_name: str, setup_id: str = "Default") -> None:
+        """Sets the scene group for a subarea
+
+        param subarea_id: the name of the subarea
+        param group_name: the name of the group
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[3] = group_name
+
+    def set_entity_group(self, subarea_id: str, group_name: str, setup_id: str = "Default") -> None:
+        """Sets the entity group for a subarea
+
+        param subarea_id: the name of the subarea
+        param group_name: the name of the group
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[4] = group_name
+
+    def set_tilegroup_group(self, subarea_id: str, group_name: str, setup_id: str = "Default") -> None:
+        """Sets the tilegroup group for a subarea
+
+        param subarea_id: the name of the subarea
+        param group_name: the name of the group
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[5] = group_name
+
+    def set_visual_preset(self, subarea_id: str, preset_name: str, setup_id: str = "Default") -> None:
+        """Sets the visual preset for a subarea
+
+        param subarea_id: the name of the subarea
+        param preset_name: the name of the preset
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[6] = preset_name
+
+    def set_sound_preset(self, subarea_id: str, preset_name: str, setup_id: str = "Default") -> None:
+        """Sets the sound preset for a subarea
+
+        param subarea_id: the name of the subarea
+        param preset_name: the name of the preset
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[7] = preset_name
+
+    def set_music_preset(self, subarea_id: str, preset_name: str, setup_id: str = "Default") -> None:
+        """Sets the music preset for a subarea
+
+        param subarea_id: the name of the subarea
+        param preset_name: the name of the preset
+        param setup_id: the name of the setup the subarea is in"""
+        self.get_subarea_config(subarea_id, setup_id).asItemsIds[8] = preset_name
+
+    def get_charclass_group(self, group_id: str) -> Container:
+        """Gets a charclass group
+
+        param group_id: the name of the group
+        returns: the charclass group"""
+        return next(group for group in self.charclass_groups if group.sId == group_id)
+
+    def add_charclass_group(self, group_id: str, charclasses: list[str] = []) -> Container:
+        """Adds a new charclass group
+
+        param group_id: the name of the new group
+        param charclasses: the charclasses in this group
+        returns: the newly created charclass group"""
+        if next((group for group in self.charclass_groups if group.sId == group_id), None):
+            raise ValueError(f"Charclass {group_id} is already present")
+
+        new_group = Container(
+            {"sId": group_id, "vsCharClassesIds": ListContainer(charclasses) if charclasses else ListContainer()}
+        )
+        self.raw.Root.pSubareaManager.vCharclassGroups.append(new_group)
 
-    def get_subarea_config(self, id: str, setup_id: str) -> Container:
-        return next(config for config in self.get_subarea_setup(setup_id).vSubareaConfigs if config.sId == id)
+        return new_group

tests/formats/test_brsa.py

@@ -30,3 +30,58 @@ def test_dread_brsa_100(dread_tree_100, brsa_path):
 @pytest.mark.parametrize("brsa_path", bossrush_assets)
 def test_dread_brsa_210(dread_tree_210, brsa_path):
     parse_build_compare_editor(Brsa, dread_tree_210, brsa_path)
+
+
+def test_add_setup(dread_tree_100):
+    subareas = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brsa", Brsa)
+    subareas.add_setup("Test")
+
+    assert len(list(subareas.subarea_setups)) == 18
+
+    with pytest.raises(ValueError, match=r"Setup .+? is already present"):
+        subareas.add_setup("Default")
+
+
+def test_add_subarea_config(dread_tree_100):
+    subareas = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brsa", Brsa)
+    new_config = subareas.add_subarea_config("camera_test")
+
+    assert len(subareas.get_subarea_setup("Default").vSubareaConfigs) == 79
+
+    with pytest.raises(ValueError, match=r"Config for .+? is already present in .+?"):
+        subareas.add_subarea_config("collision_camera_000")
+
+    subareas.set_scenario_collider("camera_test", "collider_test")
+    subareas.set_light_group("camera_test", "lg_test")
+    subareas.set_sound_group("camera_test", "ssg_test")
+    subareas.set_scene_group("camera_test", "sg_test")
+    subareas.set_entity_group("camera_test", "eg_test")
+    subareas.set_tilegroup_group("camera_test", "bg_test")
+    subareas.set_visual_preset("camera_test", "visual_test")
+    subareas.set_sound_preset("camera_test", "sound_test")
+    subareas.set_music_preset("camera_test", "music_test")
+
+    for index, to_check in enumerate(
+        [
+            "collider_test",
+            "lg_test",
+            "ssg_test",
+            "sg_test",
+            "eg_test",
+            "bg_test",
+            "visual_test",
+            "sound_test",
+            "music_test",
+        ]
+    ):
+        assert new_config.asItemsIds[index] == to_check
+
+
+def test_charclasses(dread_tree_100):
+    subareas = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brsa", Brsa)
+    subareas.add_charclass_group("Test", ["klaida"])
+
+    assert len(subareas.get_charclass_group("Test").vsCharClassesIds) == 1
+
+    with pytest.raises(ValueError, match=r"Charclass .+? is already present"):
+        subareas.add_charclass_group("No Enemies")