Clean up base defense generation.

theater/controlpoint.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import itertools
 import re
 from typing import Dict, List, TYPE_CHECKING
 from enum import Enum
@@ -50,7 +51,8 @@ class ControlPoint(MissionTarget):
         self.id = id
         self.full_name = name
         self.at = at
-        self.ground_objects: List[TheaterGroundObject] = []
+        self.connected_objectives: List[TheaterGroundObject] = []
+        self.base_defenses: List[SamGroundObject] = []
 
         self.size = size
         self.importance = importance
@@ -64,6 +66,11 @@ class ControlPoint(MissionTarget):
         self.stances: Dict[int, CombatStance] = {}
         self.airport = None
 
+    @property
+    def ground_objects(self) -> List[TheaterGroundObject]:
+        return list(
+            itertools.chain(self.connected_objectives, self.base_defenses))
+
     @classmethod
     def from_airport(cls, airport: Airport, radials: List[int], size: int, importance: float, has_frontline=True):
         assert airport
@@ -225,17 +232,6 @@ class ControlPoint(MissionTarget):
         self.base.armor = {}
 
         # Handle cyclic dependency.
-        from .start_generator import generate_airbase_defense_group
+        from .start_generator import BaseDefenseGenerator
-        base_defense_idx = 0
+        self.base_defenses = []
-        for ground_object in self.ground_objects:
+        BaseDefenseGenerator(game, self, faction_name).generate()
-            if not isinstance(ground_object, SamGroundObject):
-                continue
-            if not ground_object.airbase_group:
-                continue
-
-            # Reset in case we replace the SAM with something else.
-            ground_object.skynet_capable = False
-            ground_object.groups = []
-            generate_airbase_defense_group(base_defense_idx, ground_object,
-                                           faction_name, game)
-            base_defense_idx += 1

theater/start_generator.py

@@ -176,7 +176,7 @@ class ControlPointGroundObjectGenerator:
         return db.FACTIONS[self.faction_name]
 
     def generate(self) -> bool:
-        self.control_point.ground_objects = []
+        self.control_point.connected_objectives = []
         if self.faction.navy_generators:
             # Even airbases can generate navies if they are close enough to the
             # water. This is not controlled by the control point definition, but
@@ -215,7 +215,7 @@ class ControlPointGroundObjectGenerator:
         g.groups = []
         if group is not None:
             g.groups.append(group)
-            self.control_point.ground_objects.append(g)
+            self.control_point.connected_objectives.append(g)
 
 
 class CarrierGroundObjectGenerator(ControlPointGroundObjectGenerator):
@@ -238,7 +238,7 @@ class CarrierGroundObjectGenerator(ControlPointGroundObjectGenerator):
         g.groups = []
         if group is not None:
             g.groups.append(group)
-        self.control_point.ground_objects.append(g)
+        self.control_point.connected_objectives.append(g)
         self.control_point.name = random.choice(carrier_names)
         return True
 
@@ -263,11 +263,40 @@ class LhaGroundObjectGenerator(ControlPointGroundObjectGenerator):
         g.groups = []
         if group is not None:
             g.groups.append(group)
-        self.control_point.ground_objects.append(g)
+        self.control_point.connected_objectives.append(g)
         self.control_point.name = random.choice(lha_names)
         return True
 
 
+class BaseDefenseGenerator:
+    def __init__(self, game: Game, control_point: ControlPoint,
+                 faction_name: str) -> None:
+        self.game = game
+        self.control_point = control_point
+        self.faction_name = faction_name
+
+    def generate(self) -> None:
+        for i in range(random.randint(3, 6)):
+            self.generate_base_defense(i)
+
+    def generate_base_defense(self, index: int) -> None:
+        position = find_location(True, self.control_point.position,
+                                 self.game.theater, 800, 3200, [], True)
+        if position is None:
+            logging.error("Could not find position for "
+                          f"{self.control_point} base defense")
+            return
+
+        group_id = self.game.next_group_id()
+
+        g = SamGroundObject(namegen.random_objective_name(), group_id,
+                            position, self.control_point, for_airbase=True)
+
+        generate_airbase_defense_group(index, g, self.faction_name,
+                                       self.game)
+        self.control_point.base_defenses.append(g)
+
+
 class AirbaseGroundObjectGenerator(ControlPointGroundObjectGenerator):
     def __init__(self, game: Game, control_point: ControlPoint,
                  templates: GroundObjectTemplates) -> None:
@@ -278,7 +307,8 @@ class AirbaseGroundObjectGenerator(ControlPointGroundObjectGenerator):
         if not super().generate():
             return False
 
-        self.generate_base_defenses()
+        BaseDefenseGenerator(self.game, self.control_point,
+                             self.faction_name).generate()
         self.generate_ground_points()
 
         if self.faction.missiles:
@@ -334,7 +364,7 @@ class AirbaseGroundObjectGenerator(ControlPointGroundObjectGenerator):
                 obj_name, category, group_id, object_id, point + template_point,
                 unit["heading"], self.control_point, unit["type"])
 
-            self.control_point.ground_objects.append(g)
+            self.control_point.connected_objectives.append(g)
 
     def generate_aa_site(self) -> None:
         obj_name = namegen.random_objective_name()
@@ -354,7 +384,7 @@ class AirbaseGroundObjectGenerator(ControlPointGroundObjectGenerator):
         group = generate_anti_air_group(self.game, g, self.faction_name)
         if group is not None:
             g.groups = [group]
-        self.control_point.ground_objects.append(g)
+        self.control_point.connected_objectives.append(g)
 
     def generate_missile_sites(self) -> None:
         for i in range(self.faction.missiles_group_count):
@@ -376,29 +406,9 @@ class AirbaseGroundObjectGenerator(ControlPointGroundObjectGenerator):
         g.groups = []
         if group is not None:
             g.groups.append(group)
-            self.control_point.ground_objects.append(g)
+            self.control_point.connected_objectives.append(g)
         return
 
-    def generate_base_defenses(self) -> None:
-        for i in range(random.randint(3, 6)):
-            self.generate_base_defense(i)
-
-    def generate_base_defense(self, index: int) -> None:
-        position = find_location(True, self.control_point.position,
-                                 self.game.theater, 800, 3200, [], True)
-        if position is None:
-            logging.error("Could not find position for "
-                          f"{self.control_point} base defense")
-            return
-
-        group_id = self.game.next_group_id()
-
-        g = SamGroundObject(namegen.random_objective_name(), group_id,
-                            position, self.control_point, for_airbase=True)
-
-        generate_airbase_defense_group(index, g, self.faction_name, self.game)
-        self.control_point.ground_objects.append(g)
-
 
 class GroundObjectGenerator:
     def __init__(self, game: Game) -> None: