diff --git a/game/theater/controlpoint.py b/game/theater/controlpoint.py index 45477bc0..245746c2 100644 --- a/game/theater/controlpoint.py +++ b/game/theater/controlpoint.py @@ -31,6 +31,7 @@ from game.point_with_heading import PointWithHeading from .theatergroundobject import ( BaseDefenseGroundObject, EwrGroundObject, + FactoryGroundObject, GenericCarrierGroundObject, SamGroundObject, TheaterGroundObject, @@ -312,6 +313,16 @@ class ControlPoint(MissionTarget, ABC): connected.extend(cp.transitive_connected_friendly_points(seen)) return connected + def can_recruit_ground_units(self, game: Game) -> bool: + """Returns True if this control point is capable of recruiting ground units.""" + if not game.settings.enable_new_ground_unit_recruitment: + return True + + for tgo in self.connected_objectives: + if isinstance(tgo, FactoryGroundObject) and not tgo.is_dead: + return True + return False + @property def is_carrier(self): """ diff --git a/qt_ui/windows/basemenu/QRecruitBehaviour.py b/qt_ui/windows/basemenu/QRecruitBehaviour.py index 1e429895..67a26ca6 100644 --- a/qt_ui/windows/basemenu/QRecruitBehaviour.py +++ b/qt_ui/windows/basemenu/QRecruitBehaviour.py @@ -56,7 +56,6 @@ class QRecruitBehaviour: unit_type: Type[UnitType], layout: QLayout, row: int, - disabled: bool = False, ) -> int: exist = QGroupBox() exist.setProperty("style", "buy-box") @@ -100,19 +99,31 @@ class QRecruitBehaviour: buy = QPushButton("+") buy.setProperty("style", "btn-buy") - buy.setDisabled(disabled) + buy.setDisabled(not self.enable_purchase(unit_type)) buy.setMinimumSize(16, 16) buy.setMaximumSize(16, 16) - buy.clicked.connect(lambda: self.buy(unit_type)) + + def on_buy(): + self.buy(unit_type) + buy.setDisabled(not self.enable_purchase(unit_type)) + sell.setDisabled(not self.enable_sale(unit_type)) + + buy.clicked.connect(on_buy) buy.setSizePolicy(QSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)) sell = QPushButton("-") sell.setProperty("style", "btn-sell") - sell.setDisabled(disabled) + sell.setDisabled(not self.enable_sale(unit_type)) sell.setMinimumSize(16, 16) sell.setMaximumSize(16, 16) sell.setSizePolicy(QSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)) - sell.clicked.connect(lambda: self.sell(unit_type)) + + def on_sell(): + self.sell(unit_type) + sell.setDisabled(not self.enable_sale(unit_type)) + buy.setDisabled(not self.enable_purchase(unit_type)) + + sell.clicked.connect(on_sell) info = QGroupBox() info.setProperty("style", "buy-box") @@ -123,7 +134,6 @@ class QRecruitBehaviour: unitInfo = QPushButton("i") unitInfo.setProperty("style", "btn-info") - unitInfo.setDisabled(disabled) unitInfo.setMinimumSize(16, 16) unitInfo.setMaximumSize(16, 16) unitInfo.clicked.connect(lambda: self.info(unit_type)) @@ -169,13 +179,13 @@ class QRecruitBehaviour: GameUpdateSignal.get_instance().updateBudget(self.game_model.game) def buy(self, unit_type: Type[UnitType]): + if not self.enable_purchase(unit_type): + logging.error(f"Purchase of {unit_type.id} not allowed at {self.cp.name}") + return + price = db.PRICES[unit_type] - if self.budget >= price: - self.pending_deliveries.order({unit_type: 1}) - self.budget -= price - else: - # TODO : display modal warning - logging.info("Not enough money !") + self.pending_deliveries.order({unit_type: 1}) + self.budget -= price self._update_count_label(unit_type) self.update_available_budget() @@ -189,6 +199,13 @@ class QRecruitBehaviour: self._update_count_label(unit_type) self.update_available_budget() + def enable_purchase(self, unit_type: Type[UnitType]) -> bool: + price = db.PRICES[unit_type] + return self.budget >= price + + def enable_sale(self, unit_type: Type[UnitType]) -> bool: + return True + def info(self, unit_type): self.info_window = QUnitInfoWindow(self.game_model.game, unit_type) self.info_window.show() diff --git a/qt_ui/windows/basemenu/airfield/QAircraftRecruitmentMenu.py b/qt_ui/windows/basemenu/airfield/QAircraftRecruitmentMenu.py index 970e1c6c..97058d83 100644 --- a/qt_ui/windows/basemenu/airfield/QAircraftRecruitmentMenu.py +++ b/qt_ui/windows/basemenu/airfield/QAircraftRecruitmentMenu.py @@ -77,12 +77,7 @@ class QAircraftRecruitmentMenu(QFrame, QRecruitBehaviour): ), ) for unit_type in sorted_units: - row = self.add_purchase_row( - unit_type, - task_box_layout, - row, - disabled=not self.cp.can_operate(unit_type), - ) + row = self.add_purchase_row(unit_type, task_box_layout, row) stretch = QVBoxLayout() stretch.addStretch() task_box_layout.addLayout(stretch, row, 0) @@ -97,6 +92,20 @@ class QAircraftRecruitmentMenu(QFrame, QRecruitBehaviour): main_layout.addWidget(scroll) self.setLayout(main_layout) + def enable_purchase(self, unit_type: Type[UnitType]) -> bool: + if not issubclass(unit_type, FlyingType): + return False + if not self.cp.can_operate(unit_type): + return False + return True + + def enable_sale(self, unit_type: Type[UnitType]) -> bool: + if not issubclass(unit_type, FlyingType): + return False + if not self.cp.can_operate(unit_type): + return False + return True + def buy(self, unit_type): if self.maximum_units > 0: if self.cp.unclaimed_parking(self.game_model.game) <= 0: diff --git a/qt_ui/windows/basemenu/ground_forces/QArmorRecruitmentMenu.py b/qt_ui/windows/basemenu/ground_forces/QArmorRecruitmentMenu.py index fa7d0246..08834e5c 100644 --- a/qt_ui/windows/basemenu/ground_forces/QArmorRecruitmentMenu.py +++ b/qt_ui/windows/basemenu/ground_forces/QArmorRecruitmentMenu.py @@ -67,14 +67,8 @@ class QArmorRecruitmentMenu(QFrame, QRecruitBehaviour): main_layout.addWidget(scroll) self.setLayout(main_layout) - def sell(self, unit_type: Type[UnitType]) -> None: - if self.pending_deliveries.pending_orders(unit_type) <= 0: - QMessageBox.critical( - self, - "Could not sell ground unit", - f"Attempted to cancel order of one {unit_type.id} at {self.cp.name} " - "but no orders are pending.", - QMessageBox.Ok, - ) - return - super().sell(unit_type) + def enable_purchase(self, unit_type: Type[UnitType]) -> bool: + return self.cp.can_recruit_ground_units(self.game_model.game) + + def enable_sale(self, unit_type: Type[UnitType]) -> bool: + return self.pending_deliveries.pending_orders(unit_type) > 0