From 532515fad7a068e4accdcf08920a9ff5e2f85bfe Mon Sep 17 00:00:00 2001
From: Raffson <Raffson@users.noreply.github.com>
Date: Sat, 3 Aug 2024 16:32:38 +0200
Subject: [PATCH] Better save-compat w.r.t. changes in AC's yaml

---
 game/dcs/aircrafttype.py | 20 ++++++++++----------
 game/migrator.py         |  2 --
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/game/dcs/aircrafttype.py b/game/dcs/aircrafttype.py
index 60c8c7a2..27a8087f 100644
--- a/game/dcs/aircrafttype.py
+++ b/game/dcs/aircrafttype.py
@@ -239,6 +239,16 @@ class AircraftType(UnitType[Type[FlyingType]]):
         list
     )
 
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        # Save compat: the `name` field has been renamed `variant_id`.
+        if "name" in state:
+            state["variant_id"] = state.pop("name")
+
+        # Update any existing models with new data on load.
+        updated = AircraftType.named(state["variant_id"])
+        updated.__dict__.update(state)
+        self.__dict__.update(updated.__dict__)
+
     def __post_init__(self) -> None:
         enrich = {}
         if FlightType.SEAD_SWEEP not in self.task_priorities:
@@ -419,16 +429,6 @@ class AircraftType(UnitType[Type[FlyingType]]):
     def task_priority(self, task: FlightType) -> int:
         return self.task_priorities[task]
 
-    def __setstate__(self, state: dict[str, Any]) -> None:
-        # Save compat: the `name` field has been renamed `variant_id`.
-        if "name" in state:
-            state["variant_id"] = state.pop("name")
-
-        # Update any existing models with new data on load.
-        updated = AircraftType.named(state["variant_id"])
-        state.update(updated.__dict__)
-        self.__dict__.update(state)
-
     @staticmethod
     def _migrator() -> Dict[str, str]:
         return {"F-15E Strike Eagle (AI)": "F-15E Strike Eagle"}
diff --git a/game/migrator.py b/game/migrator.py
index 2dbc178c..c96d62a1 100644
--- a/game/migrator.py
+++ b/game/migrator.py
@@ -11,7 +11,6 @@ from game.ato.flightplans.formation import FormationLayout
 from game.ato.flightplans.waypointbuilder import WaypointBuilder
 from game.ato.packagewaypoints import PackageWaypoints
 from game.data.doctrine import MODERN_DOCTRINE, COLDWAR_DOCTRINE, WWII_DOCTRINE
-from game.dcs.aircrafttype import AircraftType
 from game.theater import ParkingType, SeasonalConditions
 
 if TYPE_CHECKING:
@@ -165,7 +164,6 @@ class Migrator:
                 try_set_attr(s, "max_size", 12)
                 try_set_attr(s, "radio_presets", {})
                 try_set_attr(s, "livery_set", [])
-                s.aircraft = AircraftType.named(s.aircraft.variant_id)
                 if isinstance(s.country, str):
                     c = country_dict.get(s.country, s.country)
                     s.country = countries_by_name[c]()