Support for find_unit and find_unit_by_id

dcs/mission.py

@@ -1844,6 +1844,52 @@ def find_group_by_id(self, group_id: int) -> Optional[Group]:
                 return g
         return None
 
+    def find_unit(self, unit_name: str, coalition: Optional[Coalition] = None) -> Optional[Unit]:
+        """Search for a unit matching a specified unit name. If a coalition is indicated,
+        the search is limited to that coalition; otherwise, it extends across all coalitions.
+
+        Args:
+            unit_name: unit's name as defined in the mission file
+            coalition: coalition to search in
+
+        Returns:
+            Unit: the unit found, otherwise None
+        """
+
+        coalitions = [coalition] if coalition is not None else self.coalition.values()
+        for k in coalitions:
+            for c in k.countries.values():
+                groups = c.plane_group + c.vehicle_group + c.ship_group + c.static_group + c.helicopter_group
+                for group in groups:
+                    for _unit in group.units:
+                        if _unit.name == unit_name:
+                            return _unit
+
+        return None
+
+    def find_unit_by_id(self, unit_id: int, coalition: Optional[Coalition] = None) -> Optional[Unit]:
+        """Search for a unit matching a specified unit identifier. If a coalition is indicated,
+        the search is limited to that coalition; otherwise, it extends across all coalitions.
+
+        Args:
+            unit_id: unit identifier assigned by the mission file
+            coalition: coalition to search in
+
+        Returns:
+            Unit: the unit found, otherwise None
+        """
+
+        coalitions = [coalition] if coalition is not None else self.coalition.values()
+        for k in coalitions:
+            for c in k.countries.values():
+                groups = c.plane_group + c.vehicle_group + c.ship_group + c.static_group + c.helicopter_group
+                for group in groups:
+                    for _unit in group.units:
+                        if _unit.id == unit_id:
+                            return _unit
+
+        return None
+
     def is_red(self, country: Country) -> bool:
         """Checks if the given country object is part o the red coalition.
 

tests/test_mission.py

@@ -1290,3 +1290,57 @@ def test_geffect(self) -> None:
         m2 = Mission()
         m2.load_file("missions/saved.g-effect-sim.miz")
         self.assertEqual(m.forced_options.geffect, m2.forced_options.geffect)
+
+    def test_find_unit(self) -> None:
+        m_name = "tests/missions/big-formation-carpet-bombing.miz"
+        m = dcs.mission.Mission()
+        m.load_file(m_name)
+
+        unit_name = "Aerial-2-1"
+        non_existing_unit_name = "Aerial-7-1"
+
+        unit_id = 2
+        non_existing_unit_id = 1234
+
+        # search by unit id
+        unit = m.find_unit_by_id(unit_id)
+        assert unit is not None
+        self.assertEqual(unit.name, unit_name)
+        self.assertEqual(unit.id, unit_id)
+
+        blue_coalition = m.coalition["blue"]
+        blue_unit = m.find_unit_by_id(unit_id, blue_coalition)
+        self.assertEqual(blue_unit, unit)
+
+        red_coalition = m.coalition["red"]
+        no_unit = m.find_unit_by_id(2, red_coalition)
+        self.assertIsNone(no_unit)
+
+        no_unit = m.find_unit_by_id(non_existing_unit_id)
+        self.assertIsNone(no_unit)
+
+        no_unit = m.find_unit_by_id(non_existing_unit_id, blue_coalition)
+        self.assertIsNone(no_unit)
+
+        no_unit = m.find_unit_by_id(non_existing_unit_id, red_coalition)
+        self.assertIsNone(no_unit)
+
+        # search by unit name
+        unit = m.find_unit(unit_name)
+        assert unit is not None
+        self.assertEqual(unit.name, unit_name)
+
+        blue_unit = m.find_unit(unit_name, blue_coalition)
+        self.assertEqual(blue_unit, unit)
+
+        no_unit = m.find_unit(unit_name, red_coalition)
+        self.assertIsNone(no_unit)
+
+        no_unit = m.find_unit(non_existing_unit_name)
+        self.assertIsNone(no_unit)
+
+        no_unit = m.find_unit(non_existing_unit_name, blue_coalition)
+        self.assertIsNone(no_unit, unit)
+
+        no_unit = m.find_unit(non_existing_unit_name, red_coalition)
+        self.assertIsNone(no_unit)