Rework local drone commands

eos/effectHandlerHelpers.py

@@ -239,7 +239,7 @@ class HandledDroneCargoList(HandledList):
         HandledList.insert(self, idx, thing)
         if thing.isInvalid:
             self.remove(thing)
-            raise(HandledListActionError)
+            raise HandledListActionError(thing)
 
 
 class HandledImplantList(HandledList):

gui/fitCommands/calc/fitAddDrone.py

@@ -1,48 +1,59 @@
 import wx
-import eos.db
 from logbook import Logger
-from eos.saveddata.drone import Drone
+
+import eos.db
+from eos.exception import HandledListActionError
+from service.fit import Fit
+from service.market import Market
+
+
 pyfalog = Logger(__name__)
 
 
 class FitAddDroneCommand(wx.Command):
-    """"
+
-    from sFit.addDrone
+    def __init__(self, fitID, droneInfo):
-    """
+        wx.Command.__init__(self, True, 'Add Drone')
-    def __init__(self, fitID, itemID, amount=1):
-        wx.Command.__init__(self, True, "Drone add")
         self.fitID = fitID
-        self.itemID = itemID
+        self.droneInfo = droneInfo
-        self.amount = amount  # add x amount. If this goes over amount, removes stack
+        self.position = None
-        self.index = None
 
     def Do(self):
-        pyfalog.debug("Adding {0} drones ({1}) to fit ID: {2}", self.amount, self.itemID, self.fitID)
+        pyfalog.debug('Doing addition of drone {} to fit {}'.format(self.droneInfo, self.fitID))
-
+        fit = Fit.getInstance().getFit(self.fitID)
-        fit = eos.db.getFit(self.fitID)
+        item = Market.getInstance().getItem(self.droneInfo.itemID, eager=("attributes", "group.category"))
-        item = eos.db.getItem(self.itemID, eager=("attributes", "group.category"))
+        # If we're not adding any active drones, check if there's an inactive stack
-
+        # with enough space for new drones and use it
-        for d in fit.drones.find(item):
+        if self.droneInfo.amountActive == 0:
-            if d is not None and d.amountActive == 0 and d.amount < max(5, fit.extraAttributes["maxActiveDrones"]):
+            for d in fit.drones.find(item):
-                drone = d
+                if (
-                break
+                    d is not None and d.amountActive == 0 and
-        else:
+                    d.amount + self.droneInfo.amount) <= max(5, fit.extraAttributes["maxActiveDrones"]
-            try:
+                ):
-                drone = Drone(item)
+                    drone = d
-            except ValueError:
+                    drone.amount += self.droneInfo.amount
-                pyfalog.warning("Invalid drone: {}", item)
+                    eos.db.commit()
-                return False
+                    self.position = fit.drones.index(drone)
-
+                    return True
-            if not drone.fits(fit):
+        # Do new stack otherwise
-                return False
+        drone = self.droneInfo.toDrone()
+        if drone is None:
+            return False
+        if not drone.fits(fit):
+            pyfalog.warning('Drone does not fit')
+            return False
+        try:
             fit.drones.append(drone)
-
+        except HandledListActionError:
-        drone.amount += self.amount
+            pyfalog.warning('Failed to append to list')
+            eos.db.commit()
+            return False
         eos.db.commit()
-        self.index = fit.drones.index(drone)
+        self.position = fit.drones.index(drone)
         return True
 
     def Undo(self):
-        from .fitRemoveDrone import FitRemoveDroneCommand  # Avoid circular import
+        pyfalog.debug('Undoing addition of drone {} to fit {}'.format(self.droneInfo, self.fitID))
-        cmd = FitRemoveDroneCommand(self.fitID, self.index, self.amount)
+        from .fitRemoveDrone import FitRemoveDroneCommand
+        cmd = FitRemoveDroneCommand(fitID=self.fitID, position=self.position, amount=self.droneInfo.amount)
         return cmd.Do()

gui/fitCommands/calc/fitRemoveDrone.py

@@ -2,65 +2,61 @@ import wx
 from logbook import Logger
 
 import eos.db
-from eos.saveddata.drone import Drone
+from eos.exception import HandledListActionError
+from gui.fitCommands.helpers import DroneInfo
 from service.fit import Fit
-from service.market import Market
 
 
 pyfalog = Logger(__name__)
 
 
 class FitRemoveDroneCommand(wx.Command):
-    """"
+
-    from sFit.addDrone
+    def __init__(self, fitID, position, amount):
-    """
+        wx.Command.__init__(self, True, 'Remove Drone')
-    def __init__(self, fitID, position, amount=1):
-        wx.Command.__init__(self, True, "Drone add")
         self.fitID = fitID
         self.position = position
         self.amountToRemove = amount
-        self.savedItemID = None
+        self.savedDroneInfo = None
-        self.savedAmount = None
-        self.savedAmountActive = None
         self.removedStack = None
 
     def Do(self):
-        pyfalog.debug("Removing {0} drones for fit ID: {1}", self.amountToRemove, self.fitID)
+        pyfalog.debug('Doing removal of {} drones at position {} from fit {}'.format(self.amountToRemove, self.position, self.fitID))
         fit = Fit.getInstance().getFit(self.fitID)
         drone = fit.drones[self.position]
-        self.savedItemID = drone.itemID
+        self.savedDroneInfo = DroneInfo.fromDrone(drone)
-        self.savedAmount = drone.amount
-        self.savedAmountActive = drone.amountActive
 
         drone.amount -= self.amountToRemove
         if drone.amountActive > 0:
             drone.amountActive -= self.amountToRemove
 
         if drone.amount == 0:
-            del fit.drones[self.position]
+            fit.drones.remove(drone)
             self.removedStack = True
         else:
             self.removedStack = False
+
         eos.db.commit()
         return True
 
     def Undo(self):
+        pyfalog.debug('Undoing removal of {} drones at position {} from fit {}'.format(self.amountToRemove, self.position, self.fitID))
         fit = Fit.getInstance().getFit(self.fitID)
         if self.removedStack:
-            droneItem = Market.getInstance().getItem(self.savedItemID, eager=("attributes", "group.category"))
+            drone = self.savedDroneInfo.toDrone()
-            try:
+            if drone is None:
-                drone = Drone(droneItem)
-            except ValueError:
-                pyfalog.warning("Invalid drone: {}", droneItem)
                 return False
             if not drone.fits(fit):
                 return False
-            drone.amount = self.savedAmount
+            try:
-            drone.amountActive = self.savedAmountActive
+                fit.drones.insert(self.position, drone)
-            fit.drones.insert(self.position, drone)
+            except HandledListActionError:
+                pyfalog.warning('Failed to insert to list')
+                eos.db.commit()
+                return False
         else:
             drone = fit.drones[self.position]
-            drone.amount = self.savedAmount
+            drone.amount = self.savedDroneInfo.amount
-            drone.amountActive = self.savedAmountActive
+            drone.amountActive = self.savedDroneInfo.amountActive
         eos.db.commit()
         return True

gui/fitCommands/guiAddDrone.py

@@ -3,6 +3,7 @@ from service.fit import Fit
 
 import gui.mainFrame
 from gui import globalEvents as GE
+from gui.fitCommands.helpers import DroneInfo
 from .calc.fitAddDrone import FitAddDroneCommand
 
 
@@ -16,7 +17,7 @@ class GuiAddDroneCommand(wx.Command):
         self.itemID = itemID
 
     def Do(self):
-        cmd = FitAddDroneCommand(self.fitID, self.itemID)
+        cmd = FitAddDroneCommand(fitID=self.fitID, droneInfo=DroneInfo(itemID=self.itemID, amount=1, amountActive=0))
         if self.internal_history.Submit(cmd):
             self.sFit.recalc(self.fitID)
             wx.PostEvent(self.mainFrame, GE.FitChanged(fitID=self.fitID))

gui/fitCommands/helpers.py

@@ -4,6 +4,7 @@ import eos.db
 from eos.const import FittingModuleState
 from eos.saveddata.booster import Booster
 from eos.saveddata.cargo import Cargo
+from eos.saveddata.drone import Drone
 from eos.saveddata.module import Module
 from service.market import Market
 from utils.repr import makeReprStr
@@ -144,6 +145,38 @@ class CargoInfo:
         return makeReprStr(self, ['itemID', 'amount'])
 
 
+class DroneInfo:
+
+    def __init__(self, itemID, amount, amountActive):
+        self.itemID = itemID
+        self.amount = amount
+        self.amountActive = amountActive
+
+    @classmethod
+    def fromDrone(cls, drone):
+        if drone is None:
+            return None
+        info = cls(
+            itemID=drone.itemID,
+            amount=drone.amount,
+            amountActive=drone.amountActive)
+        return info
+
+    def toDrone(self):
+        item = Market.getInstance().getItem(self.itemID)
+        try:
+            drone = Drone(item)
+        except ValueError:
+            pyfalog.warning('Invalid item: {}'.format(self.itemID))
+            return None
+        drone.amount = self.amount
+        drone.amountActive = self.amountActive
+        return drone
+
+    def __repr__(self):
+        return makeReprStr(self, ['itemID', 'amount', 'amountActive'])
+
+
 def stateLimit(itemIdentity):
     item = Market.getInstance().getItem(itemIdentity)
     if {'moduleBonusAssaultDamageControl', 'moduleBonusIndustrialInvulnerability'}.intersection(item.effects):