Use the actual Country type instead of the name.

We want other pieces of country information (in particular the short
names). This cleans up a lot of code anyway.

As an added bonus, this now catches squadrons that used invalid names
which would previously be passed through to pydcs and... then I don't
know what would happen.

game/ato/flight.py

@@ -44,7 +44,6 @@ class Flight(SidcDescribable, RadioFrequencyContainer, TacanContainer):
     def __init__(
         self,
         package: Package,
-        country: str,
         squadron: Squadron,
         count: int,
         flight_type: FlightType,
@@ -60,7 +59,6 @@ class Flight(SidcDescribable, RadioFrequencyContainer, TacanContainer):
     ) -> None:
         self.id = uuid.uuid4()
         self.package = package
-        self.country = country
         self.coalition = squadron.coalition
         self.squadron = squadron
         if claim_inv:

game/coalition.py

@@ -70,14 +70,6 @@ class Coalition:
             return 2
         return 1
 
-    @property
-    def country_name(self) -> str:
-        return self.faction.country
-
-    @property
-    def country_shortname(self) -> str:
-        return self.faction.country_shortname
-
     @property
     def opponent(self) -> Coalition:
         assert self._opponent is not None

game/commander/packagebuilder.py

@@ -26,13 +26,11 @@ class PackageBuilder:
         air_wing: AirWing,
         flight_db: Database[Flight],
         is_player: bool,
-        package_country: str,
         start_type: StartType,
         asap: bool,
     ) -> None:
         self.closest_airfields = closest_airfields
         self.is_player = is_player
-        self.package_country = package_country
         self.package = Package(location, flight_db, auto_asap=asap)
         self.air_wing = air_wing
         self.start_type = start_type
@@ -56,7 +54,6 @@ class PackageBuilder:
 
         flight = Flight(
             self.package,
-            self.package_country,
             squadron,
             plan.num_aircraft,
             plan.task,

game/commander/packagefulfiller.py

@@ -141,7 +141,6 @@ class PackageFulfiller:
             self.air_wing,
             self.flight_db,
             self.is_player,
-            self.coalition.country_name,
             self.default_start_type,
             mission.asap,
         )

game/dcs/countries.py

@@ -0,0 +1,9 @@
+from dcs.countries import country_dict
+from dcs.country import Country
+
+
+def country_with_name(name: str) -> Country:
+    for country in country_dict.values():
+        if country.name == name:
+            return country()
+    raise KeyError(f"No country found named {name}")

game/debriefing.py

@@ -157,8 +157,8 @@ class Debriefing:
         self.game = game
         self.unit_map = unit_map
 
-        self.player_country = game.blue.country_name
-        self.enemy_country = game.red.country_name
+        self.player_country = game.blue.faction.country.name
+        self.enemy_country = game.red.faction.country.name
 
         self.air_losses = self.dead_aircraft()
         self.ground_losses = self.dead_ground_units()

game/factions/faction.py

@@ -7,9 +7,8 @@ from functools import cached_property
 from typing import Optional, Dict, Type, List, Any, Iterator, TYPE_CHECKING
 
 import dcs
-from dcs.countries import country_dict
-from dcs.unittype import ShipType, StaticType
-from dcs.unittype import UnitType as DcsUnitType
+from dcs.country import Country
+from dcs.unittype import ShipType, StaticType, UnitType as DcsUnitType
 
 from game.armedforces.forcegroup import ForceGroup
 from game.data.building_data import (
@@ -29,6 +28,7 @@ from game.data.doctrine import (
 from game.data.groups import GroupRole
 from game.data.units import UnitClass
 from game.dcs.aircrafttype import AircraftType
+from game.dcs.countries import country_with_name
 from game.dcs.groundunittype import GroundUnitType
 from game.dcs.shipunittype import ShipUnitType
 from game.dcs.unittype import UnitType
@@ -45,7 +45,7 @@ class Faction:
     locales: Optional[List[str]]
 
     # Country used by this faction
-    country: str = field(default="")
+    country: Country
 
     # Country's short name used by this faction
     country_shortname: str = field(default="")
@@ -183,23 +183,16 @@ class Faction:
         return list(self.aircraft + self.awacs + self.tankers)
 
     @classmethod
-    def from_json(cls: Type[Faction], json: Dict[str, Any]) -> Faction:
-        faction = Faction(locales=json.get("locales"))
+    def from_dict(cls: Type[Faction], json: Dict[str, Any]) -> Faction:
+        try:
+            country = country_with_name(json["country"])
+        except KeyError as ex:
+            raise KeyError(
+                f'Faction\'s country ("{json.get("country")}") is not a valid DCS '
+                "country ID"
+            ) from ex
 
-        faction.country = json.get("country", "/")
-
-        country = None
-        for c in country_dict.values():
-            if c.name == faction.country:
-                country = c
-                break
-
-        if country is None:
-            raise AssertionError(
-                'Faction\'s country ("{}") is not a valid DCS country ID'.format(
-                    faction.country
-                )
-            )
+        faction = Faction(locales=json.get("locales"), country=country)
 
         faction.country_shortname = country.shortname
 

game/game.py

@@ -26,6 +26,7 @@ from .ato.flighttype import FlightType
 from .campaignloader import CampaignAirWingConfig
 from .coalition import Coalition
 from .db.gamedb import GameDb
+from .dcs.countries import country_with_name
 from .infos.information import Information
 from .profiling import logged_duration
 from .settings import Settings
@@ -179,13 +180,15 @@ class Game:
         Make sure the opposing factions are using different countries
         :return:
         """
+        # TODO: This should just be rejected and sent back to the user to fix.
+        # This isn't always something that the original faction can support.
         if player_faction.country == enemy_faction.country:
-            if player_faction.country == "USA":
-                enemy_faction.country = "USAF Aggressors"
-            elif player_faction.country == "Russia":
-                enemy_faction.country = "USSR"
+            if player_faction.country.name == "USA":
+                enemy_faction.country = country_with_name("USAF Aggressors")
+            elif player_faction.country.name == "Russia":
+                enemy_faction.country = country_with_name("USSR")
             else:
-                enemy_faction.country = "Russia"
+                enemy_faction.country = country_with_name("Russia")
 
     def faction_for(self, player: bool) -> Faction:
         return self.coalition_for(player).faction
@@ -196,13 +199,10 @@ class Game:
     def air_wing_for(self, player: bool) -> AirWing:
         return self.coalition_for(player).air_wing
 
-    def country_for(self, player: bool) -> str:
-        return self.coalition_for(player).country_name
-
     @property
     def neutral_country(self) -> Type[Country]:
         """Return the best fitting country that can be used as neutral faction in the generated mission"""
-        countries_in_use = [self.red.country_name, self.blue.country_name]
+        countries_in_use = {self.red.faction.country, self.blue.faction.country}
         if UnitedNationsPeacekeepers not in countries_in_use:
             return UnitedNationsPeacekeepers
         elif Switzerland.name not in countries_in_use:

game/missiongenerator/aircraft/aircraftgenerator.py

@@ -174,7 +174,6 @@ class AircraftGenerator:
             # TODO: Special flight type?
             flight = Flight(
                 Package(squadron.location, self.game.db.flights),
-                faction.country,
                 squadron,
                 1,
                 FlightType.BARCAP,

game/missiongenerator/cargoshipgenerator.py

@@ -29,12 +29,9 @@ class CargoShipGenerator:
                 self.generate_cargo_ship(ship)
 
     def generate_cargo_ship(self, ship: CargoShip) -> ShipGroup:
-        country = self.mission.country(
-            self.game.coalition_for(ship.player_owned).country_name
-        )
         waypoints = ship.route
         group = self.mission.ship_group(
-            country,
+            self.game.coalition_for(ship.player_owned).faction.country,
             ship.name,
             HandyWind,
             position=waypoints[0],

game/missiongenerator/convoygenerator.py

@@ -72,14 +72,11 @@ class ConvoyGenerator:
         units: dict[GroundUnitType, int],
         for_player: bool,
     ) -> VehicleGroup:
-        country = self.mission.country(self.game.coalition_for(for_player).country_name)
-        faction = self.game.faction_for(for_player)
-
         unit_types = list(units.items())
         main_unit_type, main_unit_count = unit_types[0]
 
         group = self.mission.vehicle_group(
-            country,
+            self.game.coalition_for(for_player).faction.country,
             name,
             main_unit_type.dcs_unit_type,
             position=position,

game/missiongenerator/flotgenerator.py

@@ -160,7 +160,7 @@ class FlotGenerator:
                 utype = AircraftType.named("MQ-9 Reaper")
 
             jtac = self.mission.flight_group(
-                country=self.mission.country(self.game.blue.country_name),
+                country=self.game.blue.faction.country,
                 name=namegen.next_jtac_name(),
                 aircraft_type=utype.dcs_unit_type,
                 position=position[0],
@@ -752,7 +752,7 @@ class FlotGenerator:
         spawn_heading = (
             self.conflict.heading.left if is_player else self.conflict.heading.right
         )
-        country = self.game.coalition_for(is_player).country_name
+        country = self.game.coalition_for(is_player).faction.country
         for group in groups:
             if group.role == CombatGroupRole.ARTILLERY:
                 distance_from_frontline = (
@@ -770,7 +770,7 @@ class FlotGenerator:
 
             g = self._generate_group(
                 is_player,
-                self.mission.country(country),
+                country,
                 group.unit_type,
                 group.size,
                 final_position,
@@ -786,7 +786,7 @@ class FlotGenerator:
                 self.gen_infantry_group_for_group(
                     g,
                     is_player,
-                    self.mission.country(country),
+                    country,
                     spawn_heading.opposite,
                 )
 

game/missiongenerator/logisticsgenerator.py

@@ -93,9 +93,8 @@ class LogisticsGenerator:
             "ctld.logisticunit"
         ):
             # Spawn logisticsunit at pickup zones
-            country = self.mission.country(self.flight.country)
             logistic_unit = self.mission.static_group(
-                country,
+                self.flight.squadron.coalition.faction.country,
                 f"{self.group.name}logistic",
                 Fortification.FARP_Ammo_Dump_Coating,
                 pickup_point,

game/missiongenerator/missiongenerator.py

@@ -133,19 +133,12 @@ class MissionGenerator:
             "neutrals", bullseye=Bullseye(Point(0, 0, self.mission.terrain)).to_pydcs()
         )
 
-        p_country = self.game.blue.country_name
-        e_country = self.game.red.country_name
-        self.mission.coalition["blue"].add_country(
-            country_dict[country_id_from_name(p_country)]()
-        )
-        self.mission.coalition["red"].add_country(
-            country_dict[country_id_from_name(e_country)]()
-        )
+        p_country = self.game.blue.faction.country
+        e_country = self.game.red.faction.country
+        self.mission.coalition["blue"].add_country(p_country)
+        self.mission.coalition["red"].add_country(e_country)
 
-        belligerents = [
-            country_id_from_name(p_country),
-            country_id_from_name(e_country),
-        ]
+        belligerents = {p_country, e_country}
         for country in country_dict.keys():
             if country not in belligerents:
                 self.mission.coalition["neutrals"].add_country(country_dict[country]())
@@ -243,19 +236,19 @@ class MissionGenerator:
         aircraft_generator.clear_parking_slots()
 
         aircraft_generator.generate_flights(
-            self.mission.country(self.game.blue.country_name),
+            self.game.blue.faction.country,
             self.game.blue.ato,
             tgo_generator.runways,
         )
         aircraft_generator.generate_flights(
-            self.mission.country(self.game.red.country_name),
+            self.game.red.faction.country,
             self.game.red.ato,
             tgo_generator.runways,
         )
         if not self.game.settings.perf_disable_idle_aircraft:
             aircraft_generator.spawn_unused_aircraft(
-                self.mission.country(self.game.blue.country_name),
-                self.mission.country(self.game.red.country_name),
+                self.game.blue.faction.country,
+                self.game.red.faction.country,
             )
 
         self.mission_data.flights = aircraft_generator.flights
@@ -285,7 +278,7 @@ class MissionGenerator:
             pos = Point(cast(float, d["x"]), cast(float, d["z"]), self.mission.terrain)
             if utype is not None and not self.game.position_culled(pos):
                 self.mission.static_group(
-                    country=self.mission.country(self.game.blue.country_name),
+                    country=self.game.blue.faction.country,
                     name="",
                     _type=utype,
                     hidden=True,

game/missiongenerator/tgogenerator.py

@@ -604,7 +604,7 @@ class HelipadGenerator:
             return
         # Note: Helipad are generated as neutral object in order not to interfer with
         # capture triggers
-        country = self.m.country(self.game.coalition_for(self.cp.captured).country_name)
+        country = self.game.coalition_for(self.cp.captured).faction.country
 
         for i, helipad in enumerate(self.cp.helipads):
             heading = helipad.heading.degrees
@@ -687,7 +687,7 @@ class TgoGenerator:
 
     def generate(self) -> None:
         for cp in self.game.theater.controlpoints:
-            country = self.m.country(self.game.coalition_for(cp.captured).country_name)
+            country = self.game.coalition_for(cp.captured).faction.country
 
             # Generate helipads
             helipad_gen = HelipadGenerator(

game/missiongenerator/visualsgenerator.py

@@ -99,7 +99,7 @@ class VisualsGenerator:
                             break
 
                         self.mission.static_group(
-                            self.mission.country(self.game.red.country_name),
+                            self.game.red.faction.country,
                             "",
                             _type=v,
                             position=pos,

game/squadrons/squadron.py

@@ -6,6 +6,7 @@ from collections.abc import Iterable
 from dataclasses import dataclass, field
 from typing import Optional, Sequence, TYPE_CHECKING
 
+from dcs.country import Country
 from faker import Faker
 
 from game.ato import Flight, FlightType, Package
@@ -27,7 +28,7 @@ if TYPE_CHECKING:
 class Squadron:
     name: str
     nickname: Optional[str]
-    country: str
+    country: Country
     role: str
     aircraft: AircraftType
     max_size: int
@@ -70,7 +71,7 @@ class Squadron:
             (
                 self.name,
                 self.nickname,
-                self.country,
+                self.country.id,
                 self.role,
                 self.aircraft,
             )
@@ -412,7 +413,6 @@ class Squadron:
 
         flight = Flight(
             package,
-            self.coalition.country_name,
             self,
             size,
             FlightType.FERRY,

game/squadrons/squadrondef.py

@@ -5,8 +5,10 @@ from pathlib import Path
 from typing import Optional, TYPE_CHECKING
 
 import yaml
+from dcs.country import Country
 
 from game.dcs.aircrafttype import AircraftType
+from game.dcs.countries import country_with_name
 from game.squadrons.operatingbases import OperatingBases
 from game.squadrons.pilot import Pilot
 
@@ -19,7 +21,7 @@ if TYPE_CHECKING:
 class SquadronDef:
     name: str
     nickname: Optional[str]
-    country: str
+    country: Country
     role: str
     aircraft: AircraftType
     livery: Optional[str]
@@ -76,7 +78,7 @@ class SquadronDef:
         return SquadronDef(
             name=data["name"],
             nickname=data.get("nickname"),
-            country=data["country"],
+            country=country_with_name(data["country"]),
             role=data["role"],
             aircraft=unit_type,
             livery=data.get("livery"),

game/squadrons/squadrondefloader.py

@@ -29,13 +29,13 @@ class SquadronDefLoader:
         squadrons: dict[AircraftType, list[SquadronDef]] = defaultdict(list)
         country = self.faction.country
         faction = self.faction
-        any_country = country.startswith("Combined Joint Task Forces ")
+        any_country = country.name.startswith("Combined Joint Task Forces ")
         for directory in self.squadron_directories():
             for path, squadron_def in self.load_squadrons_from(directory):
                 if not any_country and squadron_def.country != country:
                     logging.debug(
                         "Not using squadron for non-matching country (is "
-                        f"{squadron_def.country}, need {country}: {path}"
+                        f"{squadron_def.country.name}, need {country.name}: {path}"
                     )
                     continue
                 if squadron_def.aircraft not in faction.aircrafts:

game/transfers.py

@@ -350,7 +350,6 @@ class AirliftPlanner:
 
         flight = Flight(
             self.package,
-            self.game.country_for(squadron.player),
             squadron,
             flight_size,
             FlightType.TRANSPORT,

qt_ui/windows/mission/flight/QFlightCreator.py

@@ -40,7 +40,6 @@ class QFlightCreator(QDialog):
         self.game = game
         self.package = package
         self.custom_name_text = None
-        self.country = self.game.blue.country_name
 
         # Make dialog modal to prevent background windows to close unexpectedly.
         self.setModal(True)
@@ -183,7 +182,6 @@ class QFlightCreator(QDialog):
 
         flight = Flight(
             self.package,
-            self.country,
             squadron,
             # A bit of a hack to work around the old API. Not actually relevant because
             # the roster is passed explicitly. Needs a refactor.

tests/test_factions.py

@@ -48,7 +48,7 @@ class TestFactionLoader(unittest.TestCase):
         with (RESOURCES_DIR / "valid_faction.json").open("r") as data:
             faction = Faction.from_json(json.load(data))
 
-            self.assertEqual(faction.country, "USA")
+            self.assertEqual(faction.country.name, "USA")
             self.assertEqual(faction.name, "USA 2005")
             self.assertEqual(faction.authors, "Khopa")
             self.assertEqual(faction.description, "This is a test description")