Removed projected modules now keep their state for undoing

gui/fitCommands/calc/fitRemoveProjectedModule.py

@@ -1,6 +1,11 @@
 import wx
-import eos.db
 from logbook import Logger
+
+import eos.db
+from service.fit import Fit
+from gui.fitCommands.helpers import ModuleInfoCache
+
+
 pyfalog = Logger(__name__)
 
 
@@ -12,12 +17,21 @@ class FitRemoveProjectedModuleCommand(wx.Command):
         wx.Command.__init__(self, True)
         self.fitID = fitID
         self.position = position
-        self.removed_item = None
+        self.savedModInfo = None
 
     def Do(self):
-        pyfalog.debug("Removing ({0}) onto: {1}", self.fitID, self.position)
-        fit = eos.db.getFit(self.fitID)
-        self.removed_item = fit.projectedModules[self.position].itemID
+        pyfalog.debug("Removing ({}) onto: {}".format(self.fitID, self.position))
+        fit = Fit.getInstance().getFit(self.fitID)
+        mod = fit.projectedModules[self.position]
+        self.savedModInfo = ModuleInfoCache(
+            modPosition=self.position,
+            itemID=mod.itemID,
+            state=mod.state,
+            chargeID=mod.chargeID,
+            baseID=None,
+            mutaplasmidID=None,
+            mutations={})
+
         del fit.projectedModules[self.position]
 
         eos.db.commit()
@@ -25,6 +39,14 @@ class FitRemoveProjectedModuleCommand(wx.Command):
 
     def Undo(self):
         from gui.fitCommands.calc.fitAddProjectedModule import FitAddProjectedModuleCommand
-        cmd = FitAddProjectedModuleCommand(self.fitID, self.removed_item)
+        cmd = FitAddProjectedModuleCommand(
+            fitID=self.fitID,
+            newItemID=self.savedModInfo.itemID,
+            newBaseItemID=self.savedModInfo.baseID,
+            newMutaplasmidID=self.savedModInfo.mutaplasmidID,
+            newMutations=self.savedModInfo.mutations,
+            newState=self.savedModInfo.state,
+            newChargeID=self.savedModInfo.chargeID,
+            newPosition=self.savedModInfo.modPosition)
         cmd.Do()
         return True