Rework cargo fit commands

gui/fitCommands/calc/fitAddCargo.py

@@ -1,43 +1,45 @@
 import wx
-import eos.db
 from logbook import Logger
+
+import eos.db
+from eos.exception import HandledListActionError
 from eos.saveddata.cargo import Cargo
+from service.fit import Fit
+from service.market import Market
+
+
 pyfalog = Logger(__name__)
 
 
 class FitAddCargoCommand(wx.Command):
-    """"
-    from sFit.addCargo
-    """
-    def __init__(self, fitID, itemID, amount=1, replace=False):
-        wx.Command.__init__(self, True, "Cargo add")
+
+    def __init__(self, fitID, itemID, amount=1):
+        wx.Command.__init__(self, True, 'Add Cargo')
         self.fitID = fitID
         self.itemID = itemID
-        self.amount = amount  # add x amount. If this goes over amount, removes stack
-        self.replace = replace  # if this is false, we increment.
+        self.amount = amount
 
     def Do(self):
-        pyfalog.debug("Adding cargo {0} (x{1}) for fit {2}", self.itemID, self.amount, self.fitID)
-
-        fit = eos.db.getFit(self.fitID)
-        item = eos.db.getItem(self.itemID)
+        pyfalog.debug('Doing addition of cargo {} x{} to fit {}'.format(self.itemID, self.amount, self.fitID))
 
+        fit = Fit.getInstance().getFit(self.fitID)
+        item = Market.getInstance().getItem(self.itemID)
         cargo = next((x for x in fit.cargo if x.itemID == self.itemID), None)
-
         if cargo is None:
             cargo = Cargo(item)
-            fit.cargo.append(cargo)
-
-        if self.replace:
-            cargo.amount = self.amount
-        else:
-            cargo.amount += self.amount
-
+            try:
+                fit.cargo.append(cargo)
+            except HandledListActionError:
+                pyfalog.warning('Failed to append to list')
+                eos.db.commit()
+                return False
+        cargo.amount += self.amount
         eos.db.commit()
         return True
 
     def Undo(self):
-        from .fitRemoveCargo import FitRemoveCargoCommand  # Avoid circular import
-        cmd = FitRemoveCargoCommand(self.fitID, self.itemID, self.amount)
+        pyfalog.debug('Undoing addition of cargo {} x{} to fit {}'.format(self.itemID, self.amount, self.fitID))
+        from .fitRemoveCargo import FitRemoveCargoCommand
+        cmd = FitRemoveCargoCommand(fitID=self.fitID, itemID=self.itemID, amount=self.amount)
         cmd.Do()
         return True

gui/fitCommands/calc/fitRemoveCargo.py

@@ -1,47 +1,40 @@
 import wx
-import eos.db
 from logbook import Logger
+
+import eos.db
+from service.fit import Fit
+
+
 pyfalog = Logger(__name__)
 
 
 class FitRemoveCargoCommand(wx.Command):
-    """"
-    Fitting command that sets the amount for an item within the cargo.
 
-    from sFit.removeCargo
-    """
-    def __init__(self, fitID, itemID, amount=1, stack=False):
-        wx.Command.__init__(self, True, "Cargo remove")
+    def __init__(self, fitID, itemID, amount=1):
+        wx.Command.__init__(self, True, 'Remove Cargo')
         self.fitID = fitID
         self.itemID = itemID
-        self.stack = stack  # remove entire stack
-        self.amount = amount  # remove x amount. If this goes over amount, removes stack
-        self.old_amount = None
+        self.amount = amount  # Pass infinity if you want to remove stack
+        self.savedRemovedAmount = None
 
     def Do(self):
-        pyfalog.debug("Removing cargo {0} (x{1}) for fit {2}", self.itemID, self.amount, self.fitID)
-        fit = eos.db.getFit(self.fitID)
+        pyfalog.debug('Doing removal of cargo {} x{} from fit {}'.format(self.itemID, self.amount, self.fitID))
+        fit = Fit.getInstance().getFit(self.fitID)
         cargo = next((x for x in fit.cargo if x.itemID == self.itemID), None)
-
         if cargo is None:
             return False
 
-        self.old_amount = cargo.amount
+        self.savedRemovedAmount = min(cargo.amount, self.amount)
 
-        if self.amount >= cargo.amount:
-            self.stack = True  # set to full stack, this allows easier logic in the Undo function
-
-        if self.stack or self.amount >= cargo.amount:
+        cargo.amount -= self.savedRemovedAmount
+        if cargo.amount <= 0:
             fit.cargo.remove(cargo)
-            eos.db.commit()
-            return True
 
-        cargo.amount -= self.amount
         eos.db.commit()
         return True
 
     def Undo(self):
-        from gui.fitCommands.calc.fitAddCargo import FitAddCargoCommand  # Avoid circular import
-        cmd = FitAddCargoCommand(self.fitID, self.itemID, self.old_amount, True)
-        cmd.Do()
-        return True
+        pyfalog.debug('Undoing removal of cargo {} x{} from fit {}'.format(self.itemID, self.amount, self.fitID))
+        from gui.fitCommands.calc.fitAddCargo import FitAddCargoCommand
+        cmd = FitAddCargoCommand(fitID=self.fitID, itemID=self.itemID, amount=self.savedRemovedAmount)
+        return cmd.Do()