Store projected fit state on deletion for undo purposes

gui/fitCommands/calc/fitAddProjectedFit.py

@@ -1,6 +1,10 @@
 import wx
-import eos.db
 from logbook import Logger
+
+import eos.db
+from service.fit import Fit
+
+
 pyfalog = Logger(__name__)
 
 
@@ -8,26 +12,32 @@ class FitAddProjectedFitCommand(wx.Command):
     """"
     from sFit.project
     """
-    def __init__(self, fitID, projectedFitID):
+    def __init__(self, fitID, projectedFitID, status):
         wx.Command.__init__(self, True)
         self.fitID = fitID
         self.projectedFitID = projectedFitID
-        self.new_index = None
-        self.old_item = None
+        self.status = status
 
     def Do(self):
         pyfalog.debug("Projecting fit ({0}) onto: {1}", self.fitID, self.projectedFitID)
-        fit = eos.db.getFit(self.fitID)
-        projectedFit = eos.db.getFit(self.projectedFitID)
+        sFit = Fit.getInstance()
+        projectee = sFit.getFit(self.fitID)
+        projector = sFit.getFit(self.projectedFitID)
 
-        if projectedFit is None or projectedFit in fit.projectedFits:
+        if projector is None or projector in projectee.projectedFits:
             return False
 
-        fit.projectedFitDict[projectedFit.ID] = projectedFit
+        projectee.projectedFitDict[projector.ID] = projector
 
         # this bit is required -- see GH issue # 83
         eos.db.saveddata_session.flush()
-        eos.db.saveddata_session.refresh(projectedFit)
+        eos.db.saveddata_session.refresh(projector)
+
+        if self.status is not None:
+            projectionInfo = projector.getProjectionInfo(self.fitID)
+            if not projectionInfo:
+                return False
+            projectionInfo.active = self.status
 
         eos.db.commit()
         return True

gui/fitCommands/calc/fitRemoveProjectedFit.py

@@ -1,7 +1,11 @@
 import wx
-import eos.db
 from logbook import Logger
+
+import eos.db
+from service.fit import Fit
 from .fitRemoveProjectedModule import FitRemoveProjectedModuleCommand
+
+
 pyfalog = Logger(__name__)
 
 
@@ -15,23 +19,29 @@ class FitRemoveProjectedFitCommand(FitRemoveProjectedModuleCommand):
         wx.Command.__init__(self, True)
         self.fitID = fitID
         self.projectedFitID = projectedFitID
+        self.savedState = None
 
     def Do(self):
-        pyfalog.debug("Removing ({0}) onto: {1}", self.fitID, self.projectedFitID)
-        fit = eos.db.getFit(self.fitID)
-        projectedFit = eos.db.getFit(self.projectedFitID)
+        pyfalog.debug("Removing ({0}) onto: {1}".format(self.fitID, self.projectedFitID))
+        sFit = Fit.getInstance()
+        projectee = sFit.getFit(self.fitID)
+        projector = sFit.getFit(self.projectedFitID)
 
-        if projectedFit is None:
+        if projector is None:
             return False
 
-        del fit.projectedFitDict[projectedFit.ID]
+        projectionInfo = projector.getProjectionInfo(self.fitID)
+        if not projectionInfo:
+            return False
+
+        self.savedState = projectionInfo.active
+
+        del projectee.projectedFitDict[projector.ID]
 
         eos.db.commit()
         return True
 
     def Undo(self):
-        # todo: figure out if I need to return false here if the fit doesn't return true (means it was deleted)
         from gui.fitCommands.calc.fitAddProjectedFit import FitAddProjectedFitCommand
-        cmd = FitAddProjectedFitCommand(self.fitID, self.projectedFitID)
-        cmd.Do()
-        return True
+        cmd = FitAddProjectedFitCommand(self.fitID, self.projectedFitID, self.savedState)
+        return cmd.Do()

gui/fitCommands/guiAddProjected.py

@@ -37,7 +37,7 @@ class GuiAddProjectedCommand(wx.Command):
             else:
                 result = self.internal_history.Submit(FitAddProjectedModuleCommand(self.fitID, self.id, None, None, None, None, None, None))
         elif self.type == 'fit':
-            result = self.internal_history.Submit(FitAddProjectedFitCommand(self.fitID, self.id))
+            result = self.internal_history.Submit(FitAddProjectedFitCommand(self.fitID, self.id, None))
 
         if result:
             self.sFit.recalc(self.fitID)