From bd2ec12e0f039c9500ea0dd94e7b2e4f7d168fef Mon Sep 17 00:00:00 2001 From: Dan Albert Date: Fri, 12 May 2023 17:59:51 -0700 Subject: [PATCH] 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 | 2 -- game/coalition.py | 4 --- game/commander/packagebuilder.py | 3 --- game/commander/packagefulfiller.py | 1 - game/dcs/countries.py | 9 +++++++ game/debriefing.py | 4 +-- game/factions/faction.py | 21 ++++++++------- game/game.py | 18 ++++++------- .../aircraft/aircraftgenerator.py | 3 +-- game/missiongenerator/airsupportgenerator.py | 2 +- game/missiongenerator/cargoshipgenerator.py | 5 +--- game/missiongenerator/convoygenerator.py | 4 +-- game/missiongenerator/flotgenerator.py | 8 +++--- game/missiongenerator/logisticsgenerator.py | 3 +-- game/missiongenerator/missiongenerator.py | 27 +++++++------------ game/missiongenerator/tgogenerator.py | 4 +-- game/missiongenerator/visualsgenerator.py | 2 +- game/squadrons/squadron.py | 6 ++--- game/squadrons/squadrondef.py | 6 +++-- game/squadrons/squadrondefloader.py | 4 +-- game/transfers.py | 1 - .../windows/mission/flight/QFlightCreator.py | 2 -- tests/test_factions.py | 2 +- 23 files changed, 63 insertions(+), 78 deletions(-) create mode 100644 game/dcs/countries.py diff --git a/game/ato/flight.py b/game/ato/flight.py index e69d872c..81bcc820 100644 --- a/game/ato/flight.py +++ b/game/ato/flight.py @@ -37,7 +37,6 @@ class Flight(SidcDescribable): def __init__( self, package: Package, - country: str, squadron: Squadron, count: int, flight_type: FlightType, @@ -49,7 +48,6 @@ class Flight(SidcDescribable): ) -> None: self.id = uuid.uuid4() self.package = package - self.country = country self.coalition = squadron.coalition self.squadron = squadron self.squadron.claim_inventory(count) diff --git a/game/coalition.py b/game/coalition.py index 0617d796..9c3d4535 100644 --- a/game/coalition.py +++ b/game/coalition.py @@ -71,10 +71,6 @@ class Coalition: return 2 return 1 - @property - def country_name(self) -> str: - return self.faction.country - @property def opponent(self) -> Coalition: assert self._opponent is not None diff --git a/game/commander/packagebuilder.py b/game/commander/packagebuilder.py index e0dcb1bb..101c00d8 100644 --- a/game/commander/packagebuilder.py +++ b/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, diff --git a/game/commander/packagefulfiller.py b/game/commander/packagefulfiller.py index 69f746ba..2854d020 100644 --- a/game/commander/packagefulfiller.py +++ b/game/commander/packagefulfiller.py @@ -143,7 +143,6 @@ class PackageFulfiller: self.air_wing, self.flight_db, self.is_player, - self.coalition.country_name, self.default_start_type, mission.asap, ) diff --git a/game/dcs/countries.py b/game/dcs/countries.py new file mode 100644 index 00000000..7fe6805d --- /dev/null +++ b/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}") diff --git a/game/debriefing.py b/game/debriefing.py index a7ece486..e21f8ba6 100644 --- a/game/debriefing.py +++ b/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() diff --git a/game/factions/faction.py b/game/factions/faction.py index 8bab8211..04df97ab 100644 --- a/game/factions/faction.py +++ b/game/factions/faction.py @@ -7,7 +7,7 @@ from functools import cached_property from typing import Any, Dict, Iterator, List, Optional, TYPE_CHECKING, Type import dcs -from dcs.countries import country_dict +from dcs.country import Country from dcs.unittype import ShipType, StaticType, UnitType as DcsUnitType from game.armedforces.forcegroup import ForceGroup @@ -28,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 @@ -43,7 +44,7 @@ class Faction: locales: Optional[List[str]] # Country used by this faction - country: str = field(default="") + country: Country # Nice name of the faction name: str = field(default="") @@ -168,15 +169,15 @@ class Faction: @classmethod def from_dict(cls: Type[Faction], json: Dict[str, Any]) -> Faction: - faction = Faction(locales=json.get("locales")) + 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", "/") - if faction.country not in [c.name for c in country_dict.values()]: - raise AssertionError( - 'Faction\'s country ("{}") is not a valid DCS country ID'.format( - faction.country - ) - ) + faction = Faction(locales=json.get("locales"), country=country) faction.name = json.get("name", "") if not faction.name: diff --git a/game/game.py b/game/game.py index 8d2a59e0..a6ba9050 100644 --- a/game/game.py +++ b/game/game.py @@ -24,6 +24,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 .persistence import SaveManager from .profiling import logged_duration @@ -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: diff --git a/game/missiongenerator/aircraft/aircraftgenerator.py b/game/missiongenerator/aircraft/aircraftgenerator.py index fcbc5121..912ba4bd 100644 --- a/game/missiongenerator/aircraft/aircraftgenerator.py +++ b/game/missiongenerator/aircraft/aircraftgenerator.py @@ -17,8 +17,8 @@ from game.ato.flighttype import FlightType from game.ato.package import Package from game.ato.starttype import StartType from game.factions.faction import Faction -from game.missiongenerator.missiondata import MissionData from game.missiongenerator.lasercoderegistry import LaserCodeRegistry +from game.missiongenerator.missiondata import MissionData from game.radio.radios import RadioRegistry from game.radio.tacan import TacanRegistry from game.runways import RunwayData @@ -143,7 +143,6 @@ class AircraftGenerator: # TODO: Special flight type? flight = Flight( Package(squadron.location, self.game.db.flights), - faction.country, squadron, 1, FlightType.BARCAP, diff --git a/game/missiongenerator/airsupportgenerator.py b/game/missiongenerator/airsupportgenerator.py index 1dbc7c09..ca0128fa 100644 --- a/game/missiongenerator/airsupportgenerator.py +++ b/game/missiongenerator/airsupportgenerator.py @@ -73,7 +73,7 @@ class AirSupportGenerator: else self.conflict.red_cp ) - country = self.mission.country(self.game.blue.country_name) + country = self.game.blue.faction.country if not self.game.settings.disable_legacy_tanker: fallback_tanker_number = 0 diff --git a/game/missiongenerator/cargoshipgenerator.py b/game/missiongenerator/cargoshipgenerator.py index ec7e6577..460420e2 100644 --- a/game/missiongenerator/cargoshipgenerator.py +++ b/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], diff --git a/game/missiongenerator/convoygenerator.py b/game/missiongenerator/convoygenerator.py index a1fe60b7..e338617a 100644 --- a/game/missiongenerator/convoygenerator.py +++ b/game/missiongenerator/convoygenerator.py @@ -70,13 +70,11 @@ class ConvoyGenerator: units: dict[GroundUnitType, int], for_player: bool, ) -> VehicleGroup: - country = self.mission.country(self.game.coalition_for(for_player).country_name) - 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, diff --git a/game/missiongenerator/flotgenerator.py b/game/missiongenerator/flotgenerator.py index af7a97a3..9cab293d 100644 --- a/game/missiongenerator/flotgenerator.py +++ b/game/missiongenerator/flotgenerator.py @@ -151,7 +151,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], @@ -716,7 +716,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 = ( @@ -734,7 +734,7 @@ class FlotGenerator: g = self._generate_group( is_player, - self.mission.country(country), + country, group.unit_type, group.size, final_position, @@ -750,7 +750,7 @@ class FlotGenerator: self.gen_infantry_group_for_group( g, is_player, - self.mission.country(country), + country, spawn_heading.opposite, ) diff --git a/game/missiongenerator/logisticsgenerator.py b/game/missiongenerator/logisticsgenerator.py index 089a6bda..30c27b4f 100644 --- a/game/missiongenerator/logisticsgenerator.py +++ b/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, diff --git a/game/missiongenerator/missiongenerator.py b/game/missiongenerator/missiongenerator.py index 38deaaa8..ffa85f8f 100644 --- a/game/missiongenerator/missiongenerator.py +++ b/game/missiongenerator/missiongenerator.py @@ -135,19 +135,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]()) @@ -271,18 +264,18 @@ 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, ) 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, ) for flight in aircraft_generator.flights: @@ -314,7 +307,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, diff --git a/game/missiongenerator/tgogenerator.py b/game/missiongenerator/tgogenerator.py index e7683f26..69b15d65 100644 --- a/game/missiongenerator/tgogenerator.py +++ b/game/missiongenerator/tgogenerator.py @@ -599,7 +599,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 @@ -675,7 +675,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( diff --git a/game/missiongenerator/visualsgenerator.py b/game/missiongenerator/visualsgenerator.py index 14aa5af3..56a19708 100644 --- a/game/missiongenerator/visualsgenerator.py +++ b/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, diff --git a/game/squadrons/squadron.py b/game/squadrons/squadron.py index 21fed869..9e5c731a 100644 --- a/game/squadrons/squadron.py +++ b/game/squadrons/squadron.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from datetime import datetime from typing import Optional, Sequence, TYPE_CHECKING +from dcs.country import Country from faker import Faker from game.ato import Flight, FlightType, Package @@ -28,7 +29,7 @@ if TYPE_CHECKING: class Squadron: name: str nickname: Optional[str] - country: str + country: Country role: str aircraft: AircraftType max_size: int @@ -71,7 +72,7 @@ class Squadron: ( self.name, self.nickname, - self.country, + self.country.id, self.role, self.aircraft, ) @@ -418,7 +419,6 @@ class Squadron: flight = Flight( package, - self.coalition.country_name, self, size, FlightType.FERRY, diff --git a/game/squadrons/squadrondef.py b/game/squadrons/squadrondef.py index 1c009dab..c9a88a6d 100644 --- a/game/squadrons/squadrondef.py +++ b/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] @@ -71,7 +73,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"), diff --git a/game/squadrons/squadrondefloader.py b/game/squadrons/squadrondefloader.py index cfd1197e..1d79fb35 100644 --- a/game/squadrons/squadrondefloader.py +++ b/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: diff --git a/game/transfers.py b/game/transfers.py index 86f946d4..806a79a8 100644 --- a/game/transfers.py +++ b/game/transfers.py @@ -357,7 +357,6 @@ class AirliftPlanner: flight = Flight( self.package, - self.game.country_for(squadron.player), squadron, flight_size, FlightType.TRANSPORT, diff --git a/qt_ui/windows/mission/flight/QFlightCreator.py b/qt_ui/windows/mission/flight/QFlightCreator.py index aa657ded..5702304b 100644 --- a/qt_ui/windows/mission/flight/QFlightCreator.py +++ b/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. diff --git a/tests/test_factions.py b/tests/test_factions.py index 58e20667..96b5a353 100644 --- a/tests/test_factions.py +++ b/tests/test_factions.py @@ -41,7 +41,7 @@ class TestFactionLoader(unittest.TestCase): with (RESOURCES_DIR / "valid_faction.json").open("r") as data: faction = Faction.from_dict(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")