diff --git a/game/theater/controlpoint.py b/game/theater/controlpoint.py index f2a95200..69b9db2f 100644 --- a/game/theater/controlpoint.py +++ b/game/theater/controlpoint.py @@ -8,7 +8,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import Enum from functools import total_ordering -from typing import Any, Dict, Iterator, List, Optional, TYPE_CHECKING, Type +from typing import Any, Dict, Iterator, List, Optional, Set, TYPE_CHECKING, Type from dcs.mapping import Point from dcs.ships import ( @@ -292,6 +292,23 @@ class ControlPoint(MissionTarget, ABC): def is_global(self): return not self.connected_points + def transitive_connected_friendly_points( + self, seen: Optional[Set[ControlPoint]] = None + ) -> List[ControlPoint]: + if seen is None: + seen = {self} + + connected = [] + for cp in self.connected_points: + if cp.captured != self.captured: + continue + if cp in seen: + continue + seen.add(cp) + connected.append(cp) + connected.extend(cp.transitive_connected_friendly_points(seen)) + return connected + @property def is_carrier(self): """ diff --git a/game/theater/supplyroutes.py b/game/theater/supplyroutes.py new file mode 100644 index 00000000..0bbaaec1 --- /dev/null +++ b/game/theater/supplyroutes.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from typing import Iterator, List, Optional + +from game.theater.controlpoint import ControlPoint + + +class SupplyRoute: + def __init__(self, control_points: List[ControlPoint]) -> None: + self.control_points = control_points + + def __contains__(self, item: ControlPoint) -> bool: + return item in self.control_points + + def __iter__(self) -> Iterator[ControlPoint]: + yield from self.control_points + + @classmethod + def for_control_point(cls, control_point: ControlPoint) -> Optional[SupplyRoute]: + connected_friendly_points = control_point.transitive_connected_friendly_points() + if not connected_friendly_points: + return None + return SupplyRoute([control_point] + connected_friendly_points) diff --git a/qt_ui/windows/basemenu/NewUnitTransferDialog.py b/qt_ui/windows/basemenu/NewUnitTransferDialog.py index 4df54d31..a4c3f554 100644 --- a/qt_ui/windows/basemenu/NewUnitTransferDialog.py +++ b/qt_ui/windows/basemenu/NewUnitTransferDialog.py @@ -24,18 +24,18 @@ from dcs.task import PinpointStrike from dcs.unittype import UnitType from game import db -from game.theater import ControlPoint +from game.theater import ControlPoint, SupplyRoute from game.transfers import RoadTransferOrder from qt_ui.models import GameModel from qt_ui.widgets.QLabeledWidget import QLabeledWidget class TransferDestinationComboBox(QComboBox): - def __init__(self, game_model: GameModel, origin: ControlPoint) -> None: + def __init__(self, origin: ControlPoint) -> None: super().__init__() - for cp in game_model.game.theater.controlpoints: - if cp != origin and cp.captured and not cp.is_global: + for cp in SupplyRoute.for_control_point(origin): + if cp != origin and cp.captured: self.addItem(cp.name, cp) self.model().sort(0) self.setCurrentIndex(0) @@ -81,10 +81,10 @@ class UnitTransferList(QFrame): class TransferDestinationPanel(QVBoxLayout): - def __init__(self, label: str, origin: ControlPoint, game_model: GameModel) -> None: + def __init__(self, label: str, origin: ControlPoint) -> None: super().__init__() - self.source_combo_box = TransferDestinationComboBox(game_model, origin) + self.source_combo_box = TransferDestinationComboBox(origin) self.addLayout(QLabeledWidget(label, self.source_combo_box)) @property @@ -266,7 +266,7 @@ class NewUnitTransferDialog(QDialog): layout = QVBoxLayout() self.setLayout(layout) - self.dest_panel = TransferDestinationPanel("Destination:", origin, game_model) + self.dest_panel = TransferDestinationPanel("Destination:", origin) self.dest_panel.changed.connect(self.on_destination_changed) layout.addLayout(self.dest_panel) diff --git a/qt_ui/windows/basemenu/QBaseMenu2.py b/qt_ui/windows/basemenu/QBaseMenu2.py index ec6f03dc..8f287b6d 100644 --- a/qt_ui/windows/basemenu/QBaseMenu2.py +++ b/qt_ui/windows/basemenu/QBaseMenu2.py @@ -11,7 +11,7 @@ from PySide2.QtWidgets import ( ) from game import Game, db -from game.theater import ControlPoint, ControlPointType +from game.theater import ControlPoint, ControlPointType, SupplyRoute from gen.flights.flight import FlightType from qt_ui.dialogs import Dialog from qt_ui.models import GameModel @@ -89,7 +89,7 @@ class QBaseMenu2(QDialog): runway_attack_button.setProperty("style", "btn-danger") runway_attack_button.clicked.connect(self.new_package) - if self.cp.captured and not self.cp.is_global: + if self.cp.captured and self.has_transfer_destinations: transfer_button = QPushButton("Transfer Units") bottom_row.addWidget(transfer_button) transfer_button.clicked.connect(self.open_transfer_dialog) @@ -103,6 +103,10 @@ class QBaseMenu2(QDialog): GameUpdateSignal.get_instance().budgetupdated.connect(self.update_budget) self.setLayout(main_layout) + @property + def has_transfer_destinations(self) -> bool: + return SupplyRoute.for_control_point(self.cp) is not None + @property def can_repair_runway(self) -> bool: return self.cp.captured and self.cp.runway_can_be_repaired