Add a wrapper type for ground unit info.

game/theater/base.py

@@ -1,12 +1,10 @@
 import itertools
 import logging
-import typing
-from typing import Dict, Type
+from typing import Any
 
-from dcs.unittype import VehicleType
-
-from game.db import PRICES
 from game.dcs.aircrafttype import AircraftType
+from game.dcs.groundunittype import GroundUnitType
+from game.dcs.unittype import UnitType
 
 BASE_MAX_STRENGTH = 1
 BASE_MIN_STRENGTH = 0
@@ -14,8 +12,8 @@ BASE_MIN_STRENGTH = 0
 
 class Base:
     def __init__(self):
-        self.aircraft: Dict[AircraftType, int] = {}
-        self.armor: Dict[Type[VehicleType], int] = {}
+        self.aircraft: dict[AircraftType, int] = {}
+        self.armor: dict[GroundUnitType, int] = {}
         self.strength = 1
 
     @property
@@ -30,13 +28,10 @@ class Base:
     def total_armor_value(self) -> int:
         total = 0
         for unit_type, count in self.armor.items():
-            try:
-                total += PRICES[unit_type] * count
-            except KeyError:
-                logging.exception(f"No price found for {unit_type.id}")
+            total += unit_type.price * count
         return total
 
-    def total_units_of_type(self, unit_type: typing.Any) -> int:
+    def total_units_of_type(self, unit_type: UnitType) -> int:
         return sum(
             [
                 c
@@ -45,30 +40,25 @@ class Base:
             ]
         )
 
-    def commission_units(self, units: typing.Dict[typing.Any, int]):
+    def commission_units(self, units: dict[Any, int]):
         for unit_type, unit_count in units.items():
             if unit_count <= 0:
                 continue
 
-            target_dict: dict[typing.Any, int]
+            target_dict: dict[Any, int]
             if isinstance(unit_type, AircraftType):
                 target_dict = self.aircraft
-            elif issubclass(unit_type, VehicleType):
+            elif isinstance(unit_type, GroundUnitType):
                 target_dict = self.armor
             else:
-                logging.error(
-                    f"Unexpected unit type of {unit_type}: "
-                    f"{unit_type.__module__}.{unit_type.__name__}"
-                )
+                logging.error(f"Unexpected unit type of {unit_type}")
                 return
 
             target_dict[unit_type] = target_dict.get(unit_type, 0) + unit_count
 
-    def commit_losses(self, units_lost: typing.Dict[typing.Any, int]):
-
+    def commit_losses(self, units_lost: dict[Any, int]):
         for unit_type, count in units_lost.items():
-
-            target_dict: dict[typing.Any, int]
+            target_dict: dict[Any, int]
             if unit_type in self.aircraft:
                 target_dict = self.aircraft
             elif unit_type in self.armor:

game/theater/controlpoint.py

@@ -16,7 +16,6 @@ from typing import (
     Optional,
     Set,
     TYPE_CHECKING,
-    Type,
     Union,
     Sequence,
     Iterable,
@@ -32,7 +31,6 @@ from dcs.ships import (
 )
 from dcs.terrain.terrain import Airport, ParkingSlot
 from dcs.unit import Unit
-from dcs.unittype import VehicleType
 
 from game import db
 from game.point_with_heading import PointWithHeading
@@ -46,8 +44,8 @@ from .theatergroundobject import (
     GenericCarrierGroundObject,
     TheaterGroundObject,
 )
-from ..db import PRICES
 from ..dcs.aircrafttype import AircraftType
+from ..dcs.groundunittype import GroundUnitType
 from ..utils import nautical_miles
 from ..weather import Conditions
 
@@ -161,13 +159,13 @@ class AircraftAllocations:
 
 @dataclass(frozen=True)
 class GroundUnitAllocations:
-    present: dict[Type[VehicleType], int]
-    ordered: dict[Type[VehicleType], int]
-    transferring: dict[Type[VehicleType], int]
+    present: dict[GroundUnitType, int]
+    ordered: dict[GroundUnitType, int]
+    transferring: dict[GroundUnitType, int]
 
     @property
-    def all(self) -> dict[Type[VehicleType], int]:
-        combined: dict[Type[VehicleType], int] = defaultdict(int)
+    def all(self) -> dict[GroundUnitType, int]:
+        combined: dict[GroundUnitType, int] = defaultdict(int)
         for unit_type, count in itertools.chain(
             self.present.items(), self.ordered.items(), self.transferring.items()
         ):
@@ -178,11 +176,11 @@ class GroundUnitAllocations:
     def total_value(self) -> int:
         total: int = 0
         for unit_type, count in self.present.items():
-            total += PRICES[unit_type] * count
+            total += unit_type.price * count
         for unit_type, count in self.ordered.items():
-            total += PRICES[unit_type] * count
+            total += unit_type.price * count
         for unit_type, count in self.transferring.items():
-            total += PRICES[unit_type] * count
+            total += unit_type.price * count
 
         return total
 
@@ -697,10 +695,10 @@ class ControlPoint(MissionTarget, ABC):
     ) -> GroundUnitAllocations:
         on_order = {}
         for unit_bought, count in self.pending_unit_deliveries.units.items():
-            if type(unit_bought) == type and issubclass(unit_bought, VehicleType):
+            if isinstance(unit_bought, GroundUnitType):
                 on_order[unit_bought] = count
 
-        transferring: dict[Type[VehicleType], int] = defaultdict(int)
+        transferring: dict[GroundUnitType, int] = defaultdict(int)
         for transfer in transfers:
             if transfer.destination == self:
                 for unit_type, count in transfer.units.items():