Add a wrapper type for ground unit info.

game/factions/faction.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
+import itertools
 import logging
 from dataclasses import dataclass, field
-from typing import Optional, Dict, Type, List, Any, cast, Iterator
+from typing import Optional, Dict, Type, List, Any, Iterator
 
 import dcs
 from dcs.countries import country_dict
-from dcs.planes import plane_map
-from dcs.unittype import FlyingType, ShipType, VehicleType, UnitType
-from dcs.vehicles import Armor, Unarmed, Infantry, Artillery, AirDefence
+from dcs.unittype import ShipType, UnitType
 
 from game.data.building_data import (
     WW2_ALLIES_BUILDINGS,
@@ -24,7 +23,7 @@ from game.data.doctrine import (
 )
 from game.data.groundunitclass import GroundUnitClass
 from game.dcs.aircrafttype import AircraftType
-from pydcs_extensions.mod_units import MODDED_VEHICLES
+from game.dcs.groundunittype import GroundUnitType
 
 
 @dataclass
@@ -55,16 +54,16 @@ class Faction:
     tankers: List[AircraftType] = field(default_factory=list)
 
     # Available frontline units
-    frontline_units: List[Type[VehicleType]] = field(default_factory=list)
+    frontline_units: List[GroundUnitType] = field(default_factory=list)
 
     # Available artillery units
-    artillery_units: List[Type[VehicleType]] = field(default_factory=list)
+    artillery_units: List[GroundUnitType] = field(default_factory=list)
 
     # Infantry units used
-    infantry_units: List[Type[VehicleType]] = field(default_factory=list)
+    infantry_units: List[GroundUnitType] = field(default_factory=list)
 
     # Logistics units used
-    logistics_units: List[Type[VehicleType]] = field(default_factory=list)
+    logistics_units: List[GroundUnitType] = field(default_factory=list)
 
     # Possible SAMS site generators for this faction
     air_defenses: List[str] = field(default_factory=list)
@@ -135,15 +134,11 @@ class Faction:
     #: both will use it.
     unrestricted_satnav: bool = False
 
-    def has_access_to_unittype(self, unitclass: GroundUnitClass) -> bool:
-        has_access = False
-        for vehicle in unitclass.unit_list:
-            if vehicle in self.frontline_units:
+    def has_access_to_unittype(self, unit_class: GroundUnitClass) -> bool:
+        for vehicle in itertools.chain(self.frontline_units, self.artillery_units):
+            if vehicle.unit_class is unit_class:
                 return True
-            if vehicle in self.artillery_units:
-                return True
-
-        return has_access
+        return False
 
     @classmethod
     def from_json(cls: Type[Faction], json: Dict[str, Any]) -> Faction:
@@ -172,10 +167,18 @@ class Faction:
             set(faction.aircrafts + faction.awacs + faction.tankers)
         )
 
-        faction.frontline_units = load_all_vehicles(json.get("frontline_units", []))
-        faction.artillery_units = load_all_vehicles(json.get("artillery_units", []))
-        faction.infantry_units = load_all_vehicles(json.get("infantry_units", []))
-        faction.logistics_units = load_all_vehicles(json.get("logistics_units", []))
+        faction.frontline_units = [
+            GroundUnitType.named(n) for n in json.get("frontline_units", [])
+        ]
+        faction.artillery_units = [
+            GroundUnitType.named(n) for n in json.get("artillery_units", [])
+        ]
+        faction.infantry_units = [
+            GroundUnitType.named(n) for n in json.get("infantry_units", [])
+        ]
+        faction.logistics_units = [
+            GroundUnitType.named(n) for n in json.get("logistics_units", [])
+        ]
 
         faction.ewrs = json.get("ewrs", [])
 
@@ -242,55 +245,24 @@ class Faction:
         return faction
 
     @property
-    def ground_units(self) -> Iterator[Type[VehicleType]]:
+    def ground_units(self) -> Iterator[GroundUnitType]:
         yield from self.artillery_units
         yield from self.frontline_units
         yield from self.logistics_units
 
-
-def unit_loader(unit: str, class_repository: List[Any]) -> Optional[Type[UnitType]]:
-    """
-    Find unit by name
-    :param unit: Unit name as string
-    :param class_repository: Repository of classes (Either a module, a class, or a list of classes)
-    :return: The unit as a PyDCS type
-    """
-    if unit is None:
-        return None
-    elif unit in plane_map.keys():
-        return plane_map[unit]
-    else:
-        for mother_class in class_repository:
-            if getattr(mother_class, unit, None) is not None:
-                return getattr(mother_class, unit)
-            if type(mother_class) is list:
-                for m in mother_class:
-                    if m.__name__ == unit:
-                        return m
-        logging.error(f"FACTION ERROR : Unable to find {unit} in pydcs")
-        return None
-
-
-def load_vehicle(name: str) -> Optional[Type[VehicleType]]:
-    return cast(
-        Optional[FlyingType],
-        unit_loader(
-            name, [Infantry, Unarmed, Armor, AirDefence, Artillery, MODDED_VEHICLES]
-        ),
-    )
-
-
-def load_all_vehicles(data) -> List[Type[VehicleType]]:
-    items = []
-    for name in data:
-        item = load_vehicle(name)
-        if item is not None:
-            items.append(item)
-    return items
+    def infantry_with_class(
+        self, unit_class: GroundUnitClass
+    ) -> Iterator[GroundUnitType]:
+        for unit in self.infantry_units:
+            if unit.unit_class is unit_class:
+                yield unit
 
 
 def load_ship(name: str) -> Optional[Type[ShipType]]:
-    return cast(Optional[FlyingType], unit_loader(name, [dcs.ships]))
+    if (ship := getattr(dcs.ships, name, None)) is not None:
+        return ship
+    logging.error(f"FACTION ERROR : Unable to find {name} in dcs.ships")
+    return None
 
 
 def load_all_ships(data) -> List[Type[ShipType]]: