Refactor BRFLD

src/mercury_engine_data_structures/formats/brfld.py

@@ -2,7 +2,11 @@
 
 import functools
 import logging
-from typing import TYPE_CHECKING
+from collections.abc import Iterator
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+import construct
 
 from mercury_engine_data_structures.base_resource import BaseResource
 from mercury_engine_data_structures.formats import standard_format
@@ -16,67 +20,170 @@
 
 logger = logging.getLogger(__name__)
 
+ActorLink = str
+
+
+class ActorLayer(str, Enum):
+    ENTITIES = "rEntitiesLayer"
+    SOUNDS = "rSoundsLayer"
+    LIGHTS = "rLightsLayer"
+
 
 class Brfld(BaseResource):
     @classmethod
     @functools.lru_cache
     def construct_class(cls, target_game: Game) -> construct.Construct:
         return standard_format.game_model("CScenario", "49.0.2")
 
-    def actors_for_layer(self, name: str) -> dict:
-        return self.raw.Root.pScenario.rEntitiesLayer.dctSublayers[name].dctActors
+    def actors_for_sublayer(self, sublayer_name: str, actor_layer: ActorLayer = ActorLayer.ENTITIES) -> dict:
+        """
+        Gets the actors in a sublayer
+
+        param sublayer_name: the name of the sublayer to get the actors of
+        param actor_layer: the actor_layer the sublayer is in
+        returns: the actors in the sublayer"""
+        return self.raw.Root.pScenario[actor_layer].dctSublayers[sublayer_name].dctActors
+
+    def sublayers_for_actor_layer(self, actor_layer: ActorLayer = ActorLayer.ENTITIES) -> Iterator[str]:
+        """
+        Iterably gets the names of every sublayer in an actor layer
 
-    def all_layers(self) -> Iterator[str]:
-        yield from self.raw.Root.pScenario.rEntitiesLayer.dctSublayers.keys()
+        param actor_layer: the actor layer to get the sublayers of
+        returns: the name of each sublayer"""
+        yield from self.raw.Root.pScenario[actor_layer].dctSublayers.keys()
+
+    def all_actors_in_actor_layer(
+        self, actor_layer: ActorLayer = ActorLayer.ENTITIES
+    ) -> Iterator[tuple[str, str, construct.Container]]:
+        """
+        Iterably gets every actor in an actor layer
 
-    def all_actors(self) -> Iterator[tuple[str, str, construct.Container]]:
-        for layer_name, sublayer in self.raw.Root.pScenario.rEntitiesLayer.dctSublayers.items():
+        param actor_layer: the actor layer to get the actors of
+        returns: for each actor in the actor layer: sublayer name, actor name, actor"""
+        for sublayer_name, sublayer in self.raw.Root.pScenario[actor_layer].dctSublayers.items():
             for actor_name, actor in sublayer.dctActors.items():
-                yield layer_name, actor_name, actor
+                yield sublayer_name, actor_name, actor
 
-    def follow_link(self, link: str):
+    def follow_link(self, link: str) -> Any | None:
+        """
+        Gets the object a link is referencing
+
+        param link: the link to follow
+        returns: the part of the BRFLD link is referencing"""
         if link != "{EMPTY}":
             result = self.raw
             for part in link.split(":"):
                 result = result[part]
             return result
 
-    def link_for_actor(self, actor_name: str, layer_name: str = "default") -> str:
-        return ":".join(["Root", "pScenario", "rEntitiesLayer", "dctSublayers", layer_name, "dctActors", actor_name])
+    def link_for_actor(
+        self, actor_name: str, sublayer_name: str = "default", actor_layer: ActorLayer = ActorLayer.ENTITIES
+    ) -> ActorLink:
+        """
+        Builds a link for an actor
 
-    def all_actor_groups(self) -> Iterator[str]:
-        yield from self.raw.Root.pScenario.rEntitiesLayer.dctActorGroups.keys()
+        param actor_name: the name of the actor
+        sublayer_name: the name of the sublayer the actor is in
+        actor_layer: the actor layer the actor is in
+        returns: a string representing where in the BRFLD the actor is"""
+        return ":".join(["Root", "pScenario", actor_layer, "dctSublayers", sublayer_name, "dctActors", actor_name])
 
-    def get_actor_group(self, group_name: str) -> list[str]:
-        return self.raw.Root.pScenario.rEntitiesLayer.dctActorGroups[group_name]
+    def actor_groups_for_actor_layer(self, actor_layer: ActorLayer = ActorLayer.ENTITIES) -> Iterator[str]:
+        """
+        Iterably gets every actor group in an actor layer
 
-    def is_actor_in_group(self, group_name: str, actor_name: str, layer_name: str = "default") -> bool:
-        return self.link_for_actor(actor_name, layer_name) in self.get_actor_group(group_name)
+        param actor_layer: the actor layer to get the actor groups of
+        returns: each actor group in the actor layer"""
+        yield from self.raw.Root.pScenario[actor_layer].dctActorGroups.keys()
 
-    def add_actor_to_group(self, group_name: str, actor_name: str, layer_name: str = "default"):
-        group = self.get_actor_group(group_name)
-        actor_link = self.link_for_actor(actor_name, layer_name)
+    def get_actor_group(self, group_name: str, actor_layer: ActorLayer = ActorLayer.ENTITIES) -> list[ActorLink]:
+        """
+        Gets an actor group
+
+        param group_name: the name of the actor group
+        param actor_layer: the actor layer the actor group is in
+        returns: a list of links to actors"""
+        return self.raw.Root.pScenario[actor_layer].dctActorGroups[group_name]
+
+    def is_actor_in_group(
+        self,
+        group_name: str,
+        actor_name: str,
+        sublayer_name: str = "default",
+        actor_layer: ActorLayer = ActorLayer.ENTITIES,
+    ) -> bool:
+        """
+        Checks if an actor is in an actor group
+
+        param group_name: the name of the actor group
+        param actor_name: the name of the actor
+        param sublayer_name: the name of the sublayer the actor is in
+        param actor_layer: the actor layer the actor is in
+        returns: true if the actor is in the actor group, false otherwise"""
+        return self.link_for_actor(actor_name, sublayer_name, actor_layer) in self.get_actor_group(
+            group_name, actor_layer
+        )
+
+    def add_actor_to_group(
+        self,
+        group_name: str,
+        actor_name: str,
+        sublayer_name: str = "default",
+        actor_layer: ActorLayer = ActorLayer.ENTITIES,
+    ) -> None:
+        """
+        Adds an actor to an actor group
+
+        param group_name: the name of the actor group
+        param actor_name: the name of the actor
+        param sublayer_name: the name of the sublayer the actor is in
+        param actor_layer: the actor layer the actor is in"""
+        group = self.get_actor_group(group_name, actor_layer)
+        actor_link = self.link_for_actor(actor_name, sublayer_name, actor_layer)
         if actor_link not in group:
             group.append(actor_link)
-
-    def remove_actor_from_group(self, group_name: str, actor_name: str, layer_name: str = "default"):
-        group = self.get_actor_group(group_name)
-        actor_link = self.link_for_actor(actor_name, layer_name)
+        else:
+            raise ValueError(f"Actor {actor_link} is already in actor group {group_name}")
+
+    def remove_actor_from_group(
+        self,
+        group_name: str,
+        actor_name: str,
+        sublayer_name: str = "default",
+        actor_layer: ActorLayer = ActorLayer.ENTITIES,
+    ) -> None:
+        """
+        Removes an actor from an actor group
+
+        param group_name: the name of the actor group
+        param actor_name: the name of the actor
+        param sublayer_name: the name of the sublayer the actor is in
+        param actor_layer: the actor layer the actor is in"""
+        group = self.get_actor_group(group_name, actor_layer)
+        actor_link = self.link_for_actor(actor_name, sublayer_name, actor_layer)
         if actor_link in group:
             group.remove(actor_link)
-
-    def add_actor_to_entity_groups(self, collision_camera_name: str, actor_name: str, layer_name: str = "default"):
+        else:
+            raise ValueError(f"Actor {actor_link} is not in actor group {group_name}")
+
+    def add_actor_to_actor_groups(
+        self,
+        collision_camera_name: str,
+        actor_name: str,
+        sublayer_name: str = "default",
+        actor_layer: ActorLayer = ActorLayer.ENTITIES,
+    ) -> None:
         """
-        adds an actor to all entity groups starting with "eg_" + collision_camera_name
+        Adds an actor to all actor groups starting with collision_camera_name
 
-        param collision_camera_name: name of the collision camera group
-        (prefix "eg_" is added to find the entity groups)
-        param actor_name: name of the actor to add to the group
-        param layer_name: name of the layer the actor belongs to
+        param collision_camera_name: the name of the collision camera group
+        param actor_name: the name of the actor to add to the group
+        param sublayer_name: the name of the sublayer the actor belongs to
+        param actor_layer: the actor layer the sublayer belongs to
         """
         collision_camera_groups = [
-            group for group in self.all_actor_groups() if group.startswith(f"eg_{collision_camera_name}")
+            group for group in self.actor_groups_for_actor_layer(actor_layer) if group.startswith(collision_camera_name)
         ]
         for group in collision_camera_groups:
             logger.debug("Add actor %s to group %s", actor_name, group)
-            self.add_actor_to_group(group, actor_name, layer_name)
+            self.add_actor_to_group(group, actor_name, sublayer_name, actor_layer)

tests/formats/test_brfld.py

@@ -4,7 +4,7 @@
 from tests.test_lib import parse_build_compare_editor
 
 from mercury_engine_data_structures import dread_data
-from mercury_engine_data_structures.formats.brfld import Brfld
+from mercury_engine_data_structures.formats.brfld import ActorLayer, Brfld
 
 bossrush_assets = [
     "maps/levels/c10_samus/s201_bossrush_scorpius/s201_bossrush_scorpius.brfld",
@@ -30,3 +30,70 @@ def test_dread_brfld_100(dread_tree_100, brfld_path):
 @pytest.mark.parametrize("brfld_path", bossrush_assets)
 def test_dread_brfld_210(dread_tree_210, brfld_path):
     parse_build_compare_editor(Brfld, dread_tree_210, brfld_path)
+
+
+def test_get_actors_methods(dread_tree_100):
+    scenario = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brfld", Brfld)
+
+    actors_for_sublayer_names = []
+
+    for sublayer in scenario.sublayers_for_actor_layer():
+        actors_for_sublayer_names += scenario.actors_for_sublayer(sublayer).keys()
+
+    all_actors_in_actor_layer_names = [
+        actor_name for sublayer_name, actor_name, actor in scenario.all_actors_in_actor_layer()
+    ]
+
+    assert actors_for_sublayer_names == all_actors_in_actor_layer_names
+
+
+def test_follow_link(dread_tree_100):
+    scenario = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brfld", Brfld)
+
+    actor_link = scenario.link_for_actor("cubemap_fr.2_cave_ini", "cubes", ActorLayer.LIGHTS)
+
+    assert scenario.follow_link(actor_link).sName == "cubemap_fr.2_cave_ini"
+
+
+to_remove_from_actor_groups = [
+    ["eg_collision_camera_000_Default", "breakabletilegroup_052", "breakables", ActorLayer.ENTITIES],
+    ["ssg_collision_camera_000_Default", "Pos_C_Trees_R", "default", ActorLayer.SOUNDS],
+    ["lg_collision_camera_000", "spot_000_1", "cave_000_light", ActorLayer.LIGHTS],
+]
+
+
+@pytest.mark.parametrize(["actor_group", "actor_name", "sublayer_name", "actor_layer"], to_remove_from_actor_groups)
+def test_remove_actor_from_actor_group(dread_tree_100, actor_group, actor_name, sublayer_name, actor_layer):
+    scenario = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brfld", Brfld)
+
+    scenario.remove_actor_from_group(actor_group, actor_name, sublayer_name, actor_layer)
+    assert not scenario.is_actor_in_group(actor_group, actor_name, sublayer_name, actor_layer)
+
+
+def test_remove_actor_from_actor_group_raises_exception(dread_tree_100):
+    scenario = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brfld", Brfld)
+
+    with pytest.raises(ValueError, match=r"Actor .+ is not in actor group .+"):
+        scenario.remove_actor_from_group("eg_collision_camera_000_Default", "StartPoint0")
+
+
+to_add_to_actor_groups = [
+    ["eg_collision_camera_000_Default", "breakabletilegroup_000", "breakables", ActorLayer.ENTITIES],
+    ["ssg_collision_camera_000_Default", "Pos_C_LavaWindow_06", "default", ActorLayer.SOUNDS],
+    ["lg_collision_camera_000", "cubemap_006_1_bake", "emmy_006_light", ActorLayer.LIGHTS],
+]
+
+
+@pytest.mark.parametrize(["actor_group", "actor_name", "sublayer_name", "actor_layer"], to_add_to_actor_groups)
+def test_add_actor_to_actor_group(dread_tree_100, actor_group, actor_name, sublayer_name, actor_layer):
+    scenario = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brfld", Brfld)
+
+    scenario.add_actor_to_actor_groups(actor_group, actor_name, sublayer_name, actor_layer)
+    assert scenario.is_actor_in_group(actor_group, actor_name, sublayer_name, actor_layer)
+
+
+def test_add_actor_to_actor_group_raises_exception(dread_tree_100):
+    scenario = dread_tree_100.get_file("maps/levels/c10_samus/s010_cave/s010_cave.brfld", Brfld)
+
+    with pytest.raises(ValueError, match=r"Actor .+ is already in actor group .+"):
+        scenario.add_actor_to_group("eg_collision_camera_000_Default", "PRP_DB_CV_006")