Refactor WorldInfo class to include methods for retrieving entities and entity IDs

Author: shima004

adf_core_python/core/agent/info/world_info.py

@@ -1,5 +1,6 @@
-from typing import Any
+from typing import Any, Optional
 
+from rcrs_core.entities.entity import Entity
 from rcrs_core.worldmodel.changeSet import ChangeSet
 from rcrs_core.worldmodel.entityID import EntityID
 from rcrs_core.worldmodel.worldmodel import WorldModel
@@ -35,3 +36,40 @@ def set_change_set(self, change_set: ChangeSet) -> None:
             Change set
         """
         self._change_set = change_set
+
+    def get_entity(self, entity_id: EntityID) -> Optional[Entity]:
+        """
+        Get the entity
+
+        Parameters
+        ----------
+        entity_id : EntityID
+            Entity ID
+
+        Returns
+        -------
+        Optional[Entity]
+            Entity
+        """
+        return self._world_model.get_entity(entity_id)
+
+    def get_entity_ids_of_type(self, entity_type: type[Entity]) -> list[EntityID]:
+        """
+        Get the entity IDs of the specified type
+
+        Parameters
+        ----------
+        entity_type : type[Entity]
+            Entity type
+
+        Returns
+        -------
+        list[EntityID]
+            Entity IDs
+        """
+        entity_ids: list[EntityID] = []
+        for entity in self._world_model.get_entities():
+            if isinstance(entity, entity_type):
+                entity_ids.append(entity.get_id())
+
+        return entity_ids