Refactor groundobjectsgen.

In preparation for deduping the building TGOs.

gen/groundobjectsgen.py

@@ -1,8 +1,11 @@
+from __future__ import annotations
+
 import logging
 import random
-from typing import Dict, Iterator
+from typing import Dict, Iterator, Optional, TYPE_CHECKING
 
 from dcs import Mission
+from dcs.country import Country
 from dcs.statics import fortification_map, warehouse_map
 from dcs.task import (
     ActivateBeaconCommand,
@@ -10,21 +13,316 @@ from dcs.task import (
     EPLRS,
     OptAlarmState,
 )
-from dcs.unit import Ship, Vehicle
+from dcs.unit import Ship, Vehicle, Unit
-from dcs.unitgroup import StaticGroup
+from dcs.unitgroup import Group, ShipGroup, StaticGroup
+from dcs.unittype import StaticType, UnitType
 
 from game import db
 from game.data.building_data import FORTIFICATION_UNITS, FORTIFICATION_UNITS_ID
 from game.db import unit_type_from_name
+from theater import ControlPoint, TheaterGroundObject
+from theater.theatergroundobject import (
+    BuildingGroundObject, CarrierGroundObject,
+    GenericCarrierGroundObject,
+    LhaGroundObject, ShipGroundObject,
+)
 from .conflictgen import Conflict
-from .radios import RadioRegistry
+from .radios import RadioFrequency, RadioRegistry
 from .runways import RunwayData
-from .tacan import TacanBand, TacanRegistry
+from .tacan import TacanBand, TacanChannel, TacanRegistry
+
+if TYPE_CHECKING:
+    from game import Game
+
 
 FARP_FRONTLINE_DISTANCE = 10000
 AA_CP_MIN_DISTANCE = 40000
 
 
+class GenericGroundObjectGenerator:
+    def __init__(self, ground_object: TheaterGroundObject, country: Country,
+                 game: Game, mission: Mission) -> None:
+        self.ground_object = ground_object
+        self.country = country
+        self.game = game
+        self.m = mission
+
+    def generate(self) -> None:
+        # Only covers SAMs and missile sites now.
+        if self.game.position_culled(self.ground_object.position):
+            return
+
+        for group in self.ground_object.groups:
+            if not group.units:
+                logging.warning(f"Found empty group in {self.ground_object}")
+                continue
+
+            unit_type = unit_type_from_name(group.units[0].type)
+            if unit_type is None:
+                raise RuntimeError(
+                    f"Unrecognized unit type: {group.units[0].type}")
+
+            vg = self.m.vehicle_group(self.country, group.name, unit_type,
+                                      position=group.position,
+                                      heading=group.units[0].heading)
+            vg.units[0].name = self.m.string(group.units[0].name)
+            vg.units[0].player_can_drive = True
+            for i, u in enumerate(group.units):
+                if i > 0:
+                    vehicle = Vehicle(self.m.next_unit_id(),
+                                      self.m.string(u.name), u.type)
+                    vehicle.position.x = u.position.x
+                    vehicle.position.y = u.position.y
+                    vehicle.heading = u.heading
+                    vehicle.player_can_drive = True
+                    vg.add_unit(vehicle)
+
+            self.enable_eplrs(vg, unit_type)
+            self.set_alarm_state(vg)
+
+    @staticmethod
+    def enable_eplrs(group: Group, unit_type: UnitType) -> None:
+        if hasattr(unit_type, 'eplrs'):
+            if unit_type.eplrs:
+                group.points[0].tasks.append(EPLRS(group.id))
+
+    def set_alarm_state(self, group: Group) -> None:
+        if self.game.settings.perf_red_alert_state:
+            group.points[0].tasks.append(OptAlarmState(2))
+        else:
+            group.points[0].tasks.append(OptAlarmState(1))
+
+
+class BuildingSiteGenerator(GenericGroundObjectGenerator):
+    def __init__(self, ground_object: BuildingGroundObject, country: Country,
+                 game: Game, mission: Mission) -> None:
+        super().__init__(ground_object, country, game, mission)
+
+    def generate(self) -> None:
+        if self.game.position_culled(self.ground_object.position):
+            return
+
+        if self.ground_object.dcs_identifier in warehouse_map:
+            static_type = warehouse_map[self.ground_object.dcs_identifier]
+            self.generate_static(static_type)
+        elif self.ground_object.dcs_identifier in fortification_map:
+            static_type = fortification_map[self.ground_object.dcs_identifier]
+            self.generate_static(static_type)
+        elif self.ground_object.dcs_identifier in FORTIFICATION_UNITS_ID:
+            for f in FORTIFICATION_UNITS:
+                if f.id == self.ground_object.dcs_identifier:
+                    unit_type = f
+                    self.generate_vehicle_group(unit_type)
+                    break
+        else:
+            logging.error(
+                f"{self.ground_object.dcs_identifier} not found in static maps")
+
+    def generate_vehicle_group(self, unit_type: UnitType) -> None:
+        if not self.ground_object.is_dead:
+            self.m.vehicle_group(
+                country=self.country,
+                name=self.ground_object.string_identifier,
+                _type=unit_type,
+                position=self.ground_object.position,
+                heading=self.ground_object.heading,
+            )
+
+    def generate_static(self, static_type: StaticType) -> None:
+        self.m.static_group(
+            country=self.country,
+            name=self.ground_object.string_identifier,
+            _type=static_type,
+            position=self.ground_object.position,
+            heading=self.ground_object.heading,
+            dead=self.ground_object.is_dead,
+        )
+
+
+class GenericCarrierGenerator(GenericGroundObjectGenerator):
+    def __init__(self, ground_object: GenericCarrierGroundObject,
+                 control_point: ControlPoint, country: Country, game: Game,
+                 mission: Mission, radio_registry: RadioRegistry,
+                 tacan_registry: TacanRegistry, icls_alloc: Iterator[int],
+                 runways: Dict[str, RunwayData]) -> None:
+        super().__init__(ground_object, country, game, mission)
+        self.ground_object = ground_object
+        self.control_point = control_point
+        self.radio_registry = radio_registry
+        self.tacan_registry = tacan_registry
+        self.icls_alloc = icls_alloc
+        self.runways = runways
+
+    def generate(self) -> None:
+        # TODO: Require single group?
+        for group in self.ground_object.groups:
+            if not group.units:
+                logging.warning(
+                    f"Found empty carrier group in {self.control_point}")
+                continue
+
+            atc = self.radio_registry.alloc_uhf()
+            ship_group = self.configure_carrier(group, atc)
+            for unit in group.units[1:]:
+                ship_group.add_unit(self.create_ship(unit, atc))
+
+            tacan = self.tacan_registry.alloc_for_band(TacanBand.X)
+            tacan_callsign = self.tacan_callsign()
+            icls = next(self.icls_alloc)
+
+            brc = self.steam_into_wind(ship_group)
+            self.activate_beacons(ship_group, tacan, tacan_callsign, icls)
+            self.add_runway_data(brc or 0, atc, tacan, tacan_callsign, icls)
+
+    def get_carrier_type(self, group: Group) -> UnitType:
+        unit_type = unit_type_from_name(group.units[0].type)
+        if unit_type is None:
+            raise RuntimeError(
+                f"Unrecognized carrier name: {group.units[0].type}")
+        return unit_type
+
+    def configure_carrier(self, group: Group,
+                          atc_channel: RadioFrequency) -> ShipGroup:
+        unit_type = self.get_carrier_type(group)
+
+        ship_group = self.m.ship_group(self.country, group.name, unit_type,
+                                       position=group.position,
+                                       heading=group.units[0].heading)
+        ship_group.set_frequency(atc_channel.hertz)
+        ship_group.units[0].name = self.m.string(group.units[0].name)
+        return ship_group
+
+    def create_ship(self, unit: Unit, atc_channel: RadioFrequency) -> Ship:
+        ship = Ship(self.m.next_unit_id(),
+                    self.m.string(unit.name),
+                    unit_type_from_name(unit.type))
+        ship.position.x = unit.position.x
+        ship.position.y = unit.position.y
+        ship.heading = unit.heading
+        # TODO: Verify.
+        ship.set_frequency(atc_channel.hertz)
+        return ship
+
+    def steam_into_wind(self, group: ShipGroup) -> Optional[int]:
+        brc = self.m.weather.wind_at_ground.direction + 180
+        for attempt in range(5):
+            point = group.points[0].position.point_from_heading(
+                brc, 100000 - attempt * 20000)
+            if self.game.theater.is_in_sea(point):
+                group.add_waypoint(point)
+                return brc
+        return None
+
+    def tacan_callsign(self) -> str:
+        raise NotImplementedError
+
+    @staticmethod
+    def activate_beacons(group: ShipGroup, tacan: TacanChannel,
+                         callsign: str, icls: int) -> None:
+        group.points[0].tasks.append(ActivateBeaconCommand(
+            channel=tacan.number,
+            modechannel=tacan.band.value,
+            callsign=callsign,
+            unit_id=group.units[0].id,
+            aa=False
+        ))
+        group.points[0].tasks.append(ActivateICLSCommand(
+            icls, unit_id=group.units[0].id
+        ))
+
+    def add_runway_data(self, brc: int, atc: RadioFrequency,
+                        tacan: TacanChannel, callsign: str, icls: int) -> None:
+        # TODO: Make unit name usable.
+        # This relies on one control point mapping exactly
+        # to one LHA, carrier, or other usable "runway".
+        # This isn't wholly true, since the DD escorts of
+        # the carrier group are valid for helicopters, but
+        # they aren't exposed as such to the game. Should
+        # clean this up so that's possible. We can't use the
+        # unit name since it's an arbitrary ID.
+        self.runways[self.control_point.name] = RunwayData(
+            self.control_point.name,
+            brc,
+            "N/A",
+            atc=atc,
+            tacan=tacan,
+            tacan_callsign=callsign,
+            icls=icls,
+        )
+
+
+class CarrierGenerator(GenericCarrierGenerator):
+    def get_carrier_type(self, group: Group) -> UnitType:
+        unit_type = super().get_carrier_type(group)
+        if self.game.settings.supercarrier:
+            unit_type = db.upgrade_to_supercarrier(unit_type,
+                                                   self.control_point.name)
+        return unit_type
+
+    def tacan_callsign(self) -> str:
+        # TODO: Assign these properly.
+        return random.choice([
+            "STE",
+            "CVN",
+            "CVH",
+            "CCV",
+            "ACC",
+            "ARC",
+            "GER",
+            "ABR",
+            "LIN",
+            "TRU",
+        ])
+
+
+class LhaGenerator(GenericCarrierGenerator):
+    def tacan_callsign(self) -> str:
+        # TODO: Assign these properly.
+        return random.choice([
+            "LHD",
+            "LHA",
+            "LHB",
+            "LHC",
+            "LHD",
+            "LDS",
+        ])
+
+
+class ShipObjectGenerator(GenericGroundObjectGenerator):
+
+    def generate(self) -> None:
+        if self.game.position_culled(self.ground_object.position):
+            return
+
+        for group in self.ground_object.groups:
+            if not group.units:
+                logging.warning(f"Found empty group in {self.ground_object}")
+                continue
+
+            unit_type = unit_type_from_name(group.units[0].type)
+            if unit_type is None:
+                raise RuntimeError(
+                    f"Unrecognized unit type: {group.units[0].type}")
+
+            self.generate_group(group, unit_type)
+
+    def generate_group(self, group_def: Group, unit_type: UnitType):
+        group = self.m.ship_group(self.country, group_def.name, unit_type,
+                                  position=group_def.position,
+                                  heading=group_def.units[0].heading)
+        group.units[0].name = self.m.string(group_def.units[0].name)
+        # TODO: Skipping the first unit looks like copy pasta from the carrier.
+        for unit in group_def.units[1:]:
+            unit_type = unit_type_from_name(unit.type)
+            ship = Ship(self.m.next_unit_id(),
+                        self.m.string(unit.name), unit_type)
+            ship.position.x = unit.position.x
+            ship.position.y = unit.position.y
+            ship.heading = unit.heading
+            group.add_unit(ship)
+        self.set_alarm_state(group)
+
+
 class GroundObjectsGenerator:
     FARP_CAPACITY = 4
 
@@ -61,188 +359,34 @@ class GroundObjectsGenerator:
             )
 
     def generate(self):
-
         for cp in self.game.theater.controlpoints:
-
             if cp.captured:
-                country = self.game.player_country
+                country_name = self.game.player_country
             else:
-                country = self.game.enemy_country
+                country_name = self.game.enemy_country
-            side = self.m.country(country)
+            country = self.m.country(country_name)
 
             for ground_object in cp.ground_objects:
-                if ground_object.dcs_identifier == "AA":
+                if isinstance(ground_object, BuildingGroundObject):
-
+                    generator = BuildingSiteGenerator(ground_object, country,
-                    if self.game.position_culled(ground_object.position):
+                                                      self.game, self.m)
-                        continue
+                elif isinstance(ground_object, CarrierGroundObject):
-
+                    generator = CarrierGenerator(ground_object, cp, country,
-                    for g in ground_object.groups:
+                                                 self.game, self.m,
-                        if len(g.units) > 0:
+                                                 self.radio_registry,
-                            utype = unit_type_from_name(g.units[0].type)
+                                                 self.tacan_registry,
-
+                                                 self.icls_alloc, self.runways)
-                            if not ground_object.sea_object:
+                elif isinstance(ground_object, LhaGroundObject):
-                                vg = self.m.vehicle_group(side, g.name, utype, position=g.position, heading=g.units[0].heading)
+                    generator = CarrierGenerator(ground_object, cp, country,
-                                vg.units[0].name = self.m.string(g.units[0].name)
+                                                 self.game, self.m,
-                                vg.units[0].player_can_drive = True
+                                                 self.radio_registry,
-                                for i, u in enumerate(g.units):
+                                                 self.tacan_registry,
-                                    if i > 0:
+                                                 self.icls_alloc, self.runways)
-                                        vehicle = Vehicle(self.m.next_unit_id(), self.m.string(u.name), u.type)
+                elif isinstance(ground_object, ShipGroundObject):
-                                        vehicle.position.x = u.position.x
+                    generator = ShipObjectGenerator(ground_object, country,
-                                        vehicle.position.y = u.position.y
+                                                    self.game, self.m)
-                                        vehicle.heading = u.heading
-                                        vehicle.player_can_drive = True
-                                        vg.add_unit(vehicle)
-
-                                if hasattr(utype, 'eplrs'):
-                                    if utype.eplrs:
-                                        vg.points[0].tasks.append(EPLRS(vg.id))
                 else:
-                                vg = self.m.ship_group(side, g.name, utype, position=g.position,
+                    generator = GenericGroundObjectGenerator(ground_object,
-                                                          heading=g.units[0].heading)
+                                                             country, self.game,
-                                vg.units[0].name = self.m.string(g.units[0].name)
+                                                             self.m)
-                                for i, u in enumerate(g.units):
+                generator.generate()
-                                    utype = unit_type_from_name(u.type)
-                                    if i > 0:
-                                        ship = Ship(self.m.next_unit_id(), self.m.string(u.name), utype)
-                                        ship.position.x = u.position.x
-                                        ship.position.y = u.position.y
-                                        ship.heading = u.heading
-                                        vg.add_unit(ship)
-
-                            if self.game.settings.perf_red_alert_state:
-                                vg.points[0].tasks.append(OptAlarmState(2))
-                            else:
-                                vg.points[0].tasks.append(OptAlarmState(1))
-
-
-                elif ground_object.dcs_identifier in ["CARRIER", "LHA"]:
-                    for g in ground_object.groups:
-                        if len(g.units) > 0:
-
-                            utype = unit_type_from_name(g.units[0].type)
-                            if ground_object.dcs_identifier == "CARRIER" and self.game.settings.supercarrier == True:
-                                utype = db.upgrade_to_supercarrier(utype, cp.name)
-
-                            sg = self.m.ship_group(side, g.name, utype, position=g.position, heading=g.units[0].heading)
-                            atc_channel = self.radio_registry.alloc_uhf()
-                            sg.set_frequency(atc_channel.hertz)
-                            sg.units[0].name = self.m.string(g.units[0].name)
-
-                            for i, u in enumerate(g.units):
-                                if i > 0:
-                                    ship = Ship(self.m.next_unit_id(), self.m.string(u.name), unit_type_from_name(u.type))
-                                    ship.position.x = u.position.x
-                                    ship.position.y = u.position.y
-                                    ship.heading = u.heading
-                                    # TODO: Verify.
-                                    ship.set_frequency(atc_channel.hertz)
-                                    sg.add_unit(ship)
-
-                            # Find carrier direction (In the wind)
-                            found_carrier_destination = False
-                            attempt = 0
-                            brc = self.m.weather.wind_at_ground.direction + 180
-                            while not found_carrier_destination and attempt < 5:
-                                point = sg.points[0].position.point_from_heading(brc, 100000-attempt*20000)
-                                if self.game.theater.is_in_sea(point):
-                                    found_carrier_destination = True
-                                    sg.add_waypoint(point)
-                                else:
-                                    attempt = attempt + 1
-
-                            # Set UP TACAN and ICLS
-                            tacan = self.tacan_registry.alloc_for_band(TacanBand.X)
-                            icls_channel = next(self.icls_alloc)
-                            # TODO: Assign these properly.
-                            if ground_object.dcs_identifier == "CARRIER":
-                                tacan_callsign = random.choice([
-                                    "STE",
-                                    "CVN",
-                                    "CVH",
-                                    "CCV",
-                                    "ACC",
-                                    "ARC",
-                                    "GER",
-                                    "ABR",
-                                    "LIN",
-                                    "TRU",
-                                ])
-                            else:
-                                tacan_callsign = random.choice([
-                                    "LHD",
-                                    "LHA",
-                                    "LHB",
-                                    "LHC",
-                                    "LHD",
-                                    "LDS",
-                                ])
-                            sg.points[0].tasks.append(ActivateBeaconCommand(
-                                channel=tacan.number,
-                                modechannel=tacan.band.value,
-                                callsign=tacan_callsign,
-                                unit_id=sg.units[0].id,
-                                aa=False
-                            ))
-                            sg.points[0].tasks.append(ActivateICLSCommand(
-                                icls_channel,
-                                unit_id=sg.units[0].id
-                            ))
-                            # TODO: Make unit name usable.
-                            # This relies on one control point mapping exactly
-                            # to one LHA, carrier, or other usable "runway".
-                            # This isn't wholly true, since the DD escorts of
-                            # the carrier group are valid for helicopters, but
-                            # they aren't exposed as such to the game. Should
-                            # clean this up so that's possible. We can't use the
-                            # unit name since it's an arbitrary ID.
-                            self.runways[cp.name] = RunwayData(
-                                cp.name,
-                                brc,
-                                "N/A",
-                                atc=atc_channel,
-                                tacan=tacan,
-                                tacan_callsign=tacan_callsign,
-                                icls=icls_channel,
-                            )
-
-                else:
-
-                    if self.game.position_culled(ground_object.position):
-                        continue
-
-                    static_type = None
-                    if ground_object.dcs_identifier in warehouse_map:
-                        static_type = warehouse_map[ground_object.dcs_identifier]
-                    elif ground_object.dcs_identifier in fortification_map:
-                        static_type = fortification_map[ground_object.dcs_identifier]
-                    elif ground_object.dcs_identifier in FORTIFICATION_UNITS_ID:
-                        for f in FORTIFICATION_UNITS:
-                            if f.id == ground_object.dcs_identifier:
-                                unit_type = f
-                                break
-                    else:
-                        print("Didn't find {} in static _map(s)!".format(ground_object.dcs_identifier))
-                        continue
-
-                    if static_type is None:
-                        if not ground_object.is_dead:
-                            group = self.m.vehicle_group(
-                                country=side,
-                                name=ground_object.string_identifier,
-                                _type=unit_type,
-                                position=ground_object.position,
-                                heading=ground_object.heading,
-                            )
-                            logging.info("generated {}object identifier {} with mission id {}".format(
-                                "dead " if ground_object.is_dead else "", group.name, group.id))
-                    else:
-                        group = self.m.static_group(
-                            country=side,
-                            name=ground_object.string_identifier,
-                            _type=static_type,
-                            position=ground_object.position,
-                            heading=ground_object.heading,
-                            dead=ground_object.is_dead,
-                        )
-
-                        logging.info("generated {}object identifier {} with mission id {}".format("dead " if ground_object.is_dead else "", group.name, group.id))

gen/sam/group_generator.py

@@ -6,7 +6,12 @@ from dcs.point import PointAction
 from dcs.unit import Vehicle
 
 
-class GroupGenerator():
+# TODO: Generate a group description rather than a pydcs group.
+# It appears that all of this work gets redone at miz generation time (see
+# groundobjectsgen for an example). We can do less work and include the data we
+# care about in the format we want if we just generate our own group description
+# types rather than pydcs groups.
+class GroupGenerator:
 
     def __init__(self, game, ground_object, faction = None): # faction is not mandatory because some subclasses do not use it
         self.game = game

theater/start_generator.py

@@ -276,8 +276,7 @@ class ControlPointGroundObjectGenerator:
             return
 
         for _ in range(self.faction.navy_group_count):
-            if not self.generate_ship():
+            self.generate_ship()
-                break
 
     def generate_ship(self) -> None:
         point = find_location(False, self.control_point.position,
@@ -413,6 +412,7 @@ class GroundObjectGenerator:
                 self.game.theater.controlpoints.remove(control_point)
 
     def generate_for_control_point(self, control_point: ControlPoint) -> bool:
+        generator: ControlPointGroundObjectGenerator
         if control_point.cptype == ControlPointType.AIRCRAFT_CARRIER_GROUP:
             generator = CarrierGroundObjectGenerator(self.game, control_point,
                                                      self.templates)
@@ -488,16 +488,16 @@ def find_location(on_ground: bool, near: Point, theater: ConflictTheater,
                     break
 
         if point:
-            for other in theater.controlpoints:
+            for control_point in theater.controlpoints:
                 if is_base_defense:
                     break
-                if other.position != near:
+                if control_point.position != near:
                     if point is None:
                         break
-                    if other.position.distance_to_point(point) < 30000:
+                    if control_point.position.distance_to_point(point) < 30000:
                         point = None
                         break
-                    for ground_obj in other.ground_objects:
+                    for ground_obj in control_point.ground_objects:
                         if ground_obj.position.distance_to_point(point) < 10000:
                             point = None
                             break

theater/theatergroundobject.py

@@ -133,8 +133,12 @@ class BuildingGroundObject(TheaterGroundObject):
         )
 
 
+class GenericCarrierGroundObject(TheaterGroundObject):
+    pass
+
+
 # TODO: Why is this both a CP and a TGO?
-class CarrierGroundObject(TheaterGroundObject):
+class CarrierGroundObject(GenericCarrierGroundObject):
     def __init__(self, name: str, group_id: int,
                  control_point: ControlPoint) -> None:
         super().__init__(
@@ -152,7 +156,7 @@ class CarrierGroundObject(TheaterGroundObject):
 
 
 # TODO: Why is this both a CP and a TGO?
-class LhaGroundObject(TheaterGroundObject):
+class LhaGroundObject(GenericCarrierGroundObject):
     def __init__(self, name: str, group_id: int,
                  control_point: ControlPoint) -> None:
         super().__init__(