diff --git a/game/event/event.py b/game/event/event.py index 2eb107f4..ea7702c9 100644 --- a/game/event/event.py +++ b/game/event/event.py @@ -342,8 +342,6 @@ class UnitsDeliveryEvent: def __init__(self, control_point: ControlPoint) -> None: self.to_cp = control_point self.units: Dict[Type[UnitType], int] = {} - self.bought_units: Dict[Type[UnitType], int] = {} - self.sold_units: Dict[Type[UnitType], int] = {} def __str__(self) -> str: return "Pending delivery to {}".format(self.to_cp) @@ -372,21 +370,30 @@ class UnitsDeliveryEvent: f"Refunding {count} {unit_type.id} at {self.to_cp.name}") game.adjust_budget(price * count, player=self.to_cp.captured) + def available_next_turn(self, unit_type: Type[UnitType]) -> int: + pending_units = self.units.get(unit_type) + if pending_units is None: + pending_units = 0 + current_units = self.to_cp.base.total_units_of_type(unit_type) + return pending_units + current_units + def process(self, game: Game) -> None: + bought_units: Dict[Type[UnitType], int] = {} + sold_units: Dict[Type[UnitType], int] = {} for unit_type, count in self.units.items(): coalition = "Ally" if self.to_cp.captured else "Enemy" aircraft = unit_type.id name = self.to_cp.name if count >= 0: - self.bought_units[unit_type] = count + bought_units[unit_type] = count game.message( f"{coalition} reinforcements: {aircraft} x {count} at {name}") else: - self.sold_units[unit_type] = 0 - count + sold_units[unit_type] = -count game.message( - f"{coalition} sold: {aircraft} x {0 - count} at {name}") - self.to_cp.base.commision_units(self.bought_units) - self.to_cp.base.commit_losses(self.sold_units) + f"{coalition} sold: {aircraft} x {-count} at {name}") + self.to_cp.base.commision_units(bought_units) + self.to_cp.base.commit_losses(sold_units) self.units = {} - self.bought_units = {} - self.sold_units = {} + bought_units = {} + sold_units = {} diff --git a/qt_ui/windows/basemenu/QRecruitBehaviour.py b/qt_ui/windows/basemenu/QRecruitBehaviour.py index 7e8b06da..968d2b2a 100644 --- a/qt_ui/windows/basemenu/QRecruitBehaviour.py +++ b/qt_ui/windows/basemenu/QRecruitBehaviour.py @@ -136,7 +136,7 @@ class QRecruitBehaviour: self.update_available_budget() def sell(self, unit_type): - if self.pending_deliveries.units.get(unit_type, 0) > 0 or self.cp.base.total_units_of_type(unit_type) > 0: + if self.pending_deliveries.available_next_turn(unit_type) > 0: price = db.PRICES[unit_type] self.budget += price self.pending_deliveries.sell({unit_type: 1}) diff --git a/qt_ui/windows/basemenu/airfield/QAircraftRecruitmentMenu.py b/qt_ui/windows/basemenu/airfield/QAircraftRecruitmentMenu.py index 0f497773..4fa2530e 100644 --- a/qt_ui/windows/basemenu/airfield/QAircraftRecruitmentMenu.py +++ b/qt_ui/windows/basemenu/airfield/QAircraftRecruitmentMenu.py @@ -92,6 +92,12 @@ class QAircraftRecruitmentMenu(QFrame, QRecruitBehaviour): self, "No space for additional aircraft", f"There is no parking space left at {self.cp.name} to accommodate another plane.", QMessageBox.Ok) return + # If we change our mind about selling, we want the aircraft to be put + # back in the inventory immediately. + elif self.pending_deliveries.units.get(unit_type, 0) < 0: + global_inventory = self.game_model.game.aircraft_inventory + inventory = global_inventory.for_control_point(self.cp) + inventory.add_aircraft(unit_type, 1) super().buy(unit_type) self.hangar_status.update_label() @@ -99,13 +105,18 @@ class QAircraftRecruitmentMenu(QFrame, QRecruitBehaviour): def sell(self, unit_type: UnitType): # Don't need to remove aircraft from the inventory if we're canceling # orders. - if self.pending_deliveries.units.get(unit_type, 0) <= 0 - self.cp.base.total_units_of_type(unit_type): - QMessageBox.critical( - self, "Could not sell aircraft", - f"Attempted to sell one {unit_type.id} at {self.cp.name} " - "but none are available. Are all aircraft currently " - "assigned to a mission?", QMessageBox.Ok) - return + if self.pending_deliveries.units.get(unit_type, 0) <= 0: + global_inventory = self.game_model.game.aircraft_inventory + inventory = global_inventory.for_control_point(self.cp) + try: + inventory.remove_aircraft(unit_type, 1) + except ValueError: + QMessageBox.critical( + self, "Could not sell aircraft", + f"Attempted to sell one {unit_type.id} at {self.cp.name} " + "but none are available. Are all aircraft currently " + "assigned to a mission?", QMessageBox.Ok) + return super().sell(unit_type) self.hangar_status.update_label() diff --git a/qt_ui/windows/basemenu/ground_forces/QArmorRecruitmentMenu.py b/qt_ui/windows/basemenu/ground_forces/QArmorRecruitmentMenu.py index c04cf4cc..14801ba4 100644 --- a/qt_ui/windows/basemenu/ground_forces/QArmorRecruitmentMenu.py +++ b/qt_ui/windows/basemenu/ground_forces/QArmorRecruitmentMenu.py @@ -61,13 +61,10 @@ class QArmorRecruitmentMenu(QFrame, QRecruitBehaviour): self.setLayout(main_layout) def sell(self, unit_type: UnitType): - # Don't need to remove aircraft from the inventory if we're canceling - # orders. - if self.pending_deliveries.units.get(unit_type, 0) <= 0 - self.cp.base.total_units_of_type(unit_type): + if self.pending_deliveries.available_next_turn(unit_type) <= 0: QMessageBox.critical( self, "Could not sell ground unit", f"Attempted to sell one {unit_type.id} at {self.cp.name} " "but none are available.", QMessageBox.Ok) return - super().sell(unit_type) - self.hangar_status.update_label() \ No newline at end of file + super().sell(unit_type) \ No newline at end of file