Rework implant fit commands

2019-04-13 01:56:17 +03:00
parent 6174bf8a0b
commit c997661f3a
6 changed files with 184 additions and 116 deletions
--- a/eos/effectHandlerHelpers.py
+++ b/eos/effectHandlerHelpers.py
@@ -249,23 +249,39 @@ class HandledImplantList(HandledList):
            HandledList.append(self, implant)
            self.remove(implant)
            raise HandledListActionError(implant)
+        if self.__slotCheck(implant):
+            HandledList.append(self, implant)
+            self.remove(implant)
+            raise HandledListActionError(implant)
+        HandledList.append(self, implant)

-        self.makeRoom(implant)
+    def insert(self, idx, implant):
+        if implant.isInvalid:
+            HandledList.insert(self, idx, implant)
+            self.remove(implant)
+            raise HandledListActionError(implant)
+        if self.__slotCheck(implant):
+            HandledList.insert(self, idx, implant)
+            self.remove(implant)
+            raise HandledListActionError(implant)
        HandledList.append(self, implant)

    def makeRoom(self, implant):
        # if needed, remove booster that was occupying slot
-        oldObj = next((m for m in self if m.slot == implant.slot), None)
-        if oldObj:
-            pyfalog.info("Slot {0} occupied with {1}, replacing with {2}", implant.slot, oldObj.item.name,
-                         implant.item.name)
-            itemID = oldObj.itemID
-            state = oldObj.active
+        oldObj = next((i for i in self if i.slot == implant.slot), None)
+        if oldObj is not None:
+            pyfalog.info("Slot {0} occupied with {1}, replacing with {2}", implant.slot, oldObj.item.name, implant.item.name)
+            position = self.index(oldObj)
+            from gui.fitCommands.helpers import ImplantInfo
+            implantInfo = ImplantInfo.fromImplant(oldObj)
            oldObj.itemID = 0  # hack to remove from DB. See GH issue #324
            self.remove(oldObj)
-            return itemID, state
+            return position, implantInfo
        return None, None

+    def __slotCheck(self, implant):
+        return any(i.slot == implant.slot for i in self)
+

 class HandledBoosterList(HandledList):

@@ -274,12 +290,10 @@ class HandledBoosterList(HandledList):
            HandledList.append(self, booster)
            self.remove(booster)
            raise HandledListActionError(booster)
-
        if self.__slotCheck(booster):
            HandledList.append(self, booster)
            self.remove(booster)
            raise HandledListActionError(booster)
-
        HandledList.append(self, booster)

    def insert(self, idx, booster):
@@ -287,17 +301,15 @@ class HandledBoosterList(HandledList):
            HandledList.insert(self, idx, booster)
            self.remove(booster)
            raise HandledListActionError(booster)
-
        if self.__slotCheck(booster):
            HandledList.insert(self, idx, booster)
            self.remove(booster)
            raise HandledListActionError(booster)
-
        HandledList.append(self, booster)

    def makeRoom(self, booster):
        # if needed, remove booster that was occupying slot
-        oldObj = next((m for m in self if m.slot == booster.slot), None)
+        oldObj = next((b for b in self if b.slot == booster.slot), None)
        if oldObj is not None:
            pyfalog.info("Slot {0} occupied with {1}, replacing with {2}", booster.slot, oldObj.item.name, booster.item.name)
            position = self.index(oldObj)
--- a/gui/fitCommands/calc/fitAddImplant.py
+++ b/gui/fitCommands/calc/fitAddImplant.py
@@ -1,51 +1,73 @@
 import wx
-import eos.db
 from logbook import Logger
-from eos.saveddata.implant import Implant
+
+import eos.db
+from eos.exception import HandledListActionError
+from service.fit import Fit
+
+
 pyfalog = Logger(__name__)


 class FitAddImplantCommand(wx.Command):
-    """"
-    from sFit.addImplant
-    """
-    def __init__(self, fitID, itemID, state):
-        wx.Command.__init__(self, True, "Add Implant")
+
+    def __init__(self, fitID, implantInfo, position=None):
+        wx.Command.__init__(self, True, 'Add Implant')
        self.fitID = fitID
-        self.newItemID = itemID
-        self.newState = state
-        self.newIndex = None
-        self.oldItemID = None
-        self.oldState = None
+        self.newImplantInfo = implantInfo
+        self.newPosition = position
+        self.oldImplantInfo = None
+        self.oldPosition = None

    def Do(self):
-        pyfalog.debug("Adding implant to fit ({0}) for item ID: {1}", self.fitID, self.newItemID)
+        pyfalog.debug('Doing addition of implant {} to fit {}'.format(self.newImplantInfo, self.fitID))
+        fit = Fit.getInstance().getFit(self.fitID)

-        fit = eos.db.getFit(self.fitID)
-        item = eos.db.getItem(self.newItemID, eager="attributes")
-
-        if next((x for x in fit.implants if x.itemID == self.newItemID), None):
-            return False  # already have item in list of implants
-
-        try:
-            implant = Implant(item)
-        except ValueError:
-            pyfalog.warning("Invalid item: {0}", self.newItemID)
+        if any(self.newImplantInfo.itemID == i.itemID for i in fit.implants):
+            pyfalog.debug('Skipping as such implant is already on the fit')
            return False
-        implant.active = self.newState

-        self.oldItemID, self.oldState = fit.implants.makeRoom(implant)
-        fit.implants.append(implant)
-        self.newIndex = fit.implants.index(implant)
+        newImplant = self.newImplantInfo.toImplant()
+        if newImplant is None:
+            return False
+
+        self.oldPosition, self.oldImplantInfo = fit.implants.makeRoom(newImplant)
+
+        if self.newPosition is not None:
+            try:
+                fit.implants.insert(self.newPosition, newImplant)
+            except HandledListActionError:
+                pyfalog.warning('Failed to insert to list')
+                cmd = FitAddImplantCommand(
+                    fitID=self.fitID,
+                    implantInfo=self.oldImplantInfo,
+                    position=self.oldPosition)
+                cmd.Do()
+                return False
+        else:
+            try:
+                fit.implants.append(newImplant)
+            except HandledListActionError:
+                pyfalog.warning('Failed to append to list')
+                cmd = FitAddImplantCommand(
+                    fitID=self.fitID,
+                    implantInfo=self.oldImplantInfo,
+                    position=self.oldPosition)
+                cmd.Do()
+                return False
+            self.newPosition = fit.implants.index(newImplant)
+
+        eos.db.commit()
        return True

    def Undo(self):
-        if self.oldItemID:
-            # If we had an item in the slot previously, add it back.
-            cmd = FitAddImplantCommand(self.fitID, self.oldItemID, self.oldState)
+        pyfalog.debug('Undo addition of implant {} to fit {}'.format(self.newImplantInfo, self.fitID))
+        if self.oldImplantInfo and self.oldPosition:
+            cmd = FitAddImplantCommand(
+                fitID=self.fitID,
+                implantInfo=self.oldImplantInfo,
+                position=self.oldPosition)
            return cmd.Do()
-
-        from .fitRemoveImplant import FitRemoveImplantCommand  # Avoid circular import
-        cmd = FitRemoveImplantCommand(self.fitID, self.newIndex)
-        cmd.Do()
-        return True
+        from .fitRemoveImplant import FitRemoveImplantCommand
+        cmd = FitRemoveImplantCommand(self.fitID, self.newPosition)
+        return cmd.Do()
--- a/gui/fitCommands/calc/fitRemoveBooster.py
+++ b/gui/fitCommands/calc/fitRemoveBooster.py
@@ -19,7 +19,6 @@ class FitRemoveBoosterCommand(wx.Command):

    def Do(self):
        pyfalog.debug('Doing removal of booster from position {} on fit {}'.format(self.position, self.fitID))
-
        fit = Fit.getInstance().getFit(self.fitID)
        booster = fit.boosters[self.position]
        self.savedBoosterInfo = BoosterInfo.fromBooster(booster)
@@ -29,7 +28,7 @@ class FitRemoveBoosterCommand(wx.Command):

    def Undo(self):
        pyfalog.debug('Undoing removal of booster {} on fit {}'.format(self.savedBoosterInfo, self.fitID))
-        from .fitAddBooster import FitAddBoosterCommand  # Avoid circular import
+        from .fitAddBooster import FitAddBoosterCommand
        cmd = FitAddBoosterCommand(
            fitID=self.fitID,
            boosterInfo=self.savedBoosterInfo,
--- a/gui/fitCommands/calc/fitRemoveImplant.py
+++ b/gui/fitCommands/calc/fitRemoveImplant.py
@@ -1,34 +1,36 @@
 import wx
-import eos.db
 from logbook import Logger
+
+import eos.db
+from gui.fitCommands.helpers import ImplantInfo
+from service.fit import Fit
+
+
 pyfalog = Logger(__name__)


 class FitRemoveImplantCommand(wx.Command):
-    """"
-    Fitting command that sets the amount for an item within the cargo.

-    from sFit.removeImplant
-    """
    def __init__(self, fitID, position):
-        wx.Command.__init__(self, True, "Implant remove")
+        wx.Command.__init__(self, True, 'Remove Implant')
        self.fitID = fitID
        self.position = position
-        self.savedItemID = None
-        self.savedState = None
+        self.savedImplantInfo = None

    def Do(self):
-        pyfalog.debug("Removing implant from position ({0}) for fit ID: {1}", self.position, self.fitID)
-
-        fit = eos.db.getFit(self.fitID)
+        pyfalog.debug('Doing removal of implant from position {} on fit {}'.format(self.position, self.fitID))
+        fit = Fit.getInstance().getFit(self.fitID)
        implant = fit.implants[self.position]
-        self.savedItemID = implant.itemID
-        self.savedState = implant.active
+        self.savedImplantInfo = ImplantInfo.fromImplant(implant)
        fit.implants.remove(implant)
+        eos.db.commit()
        return True

    def Undo(self):
-        from gui.fitCommands.calc.fitAddImplant import FitAddImplantCommand  # Avoid circular import
-        cmd = FitAddImplantCommand(self.fitID, self.savedItemID, self.savedState)
-        cmd.Do()
-        return True
+        pyfalog.debug('Undoing removal of implant {} on fit {}'.format(self.savedImplantInfo, self.fitID))
+        from .fitAddImplant import FitAddImplantCommand
+        cmd = FitAddImplantCommand(
+            fitID=self.fitID,
+            boosterInfo=self.savedBoosterInfo,
+            position=self.position)
+        return cmd.Do()
--- a/gui/fitCommands/guiAddImplant.py
+++ b/gui/fitCommands/guiAddImplant.py
@@ -2,8 +2,9 @@ import wx
 from service.fit import Fit

 import gui.mainFrame
-from gui import globalEvents as GE
 from eos.const import ImplantLocation
+from gui import globalEvents as GE
+from gui.fitCommands.helpers import ImplantInfo
 from .calc.fitAddImplant import FitAddImplantCommand
 from .calc.fitChangeImplantLocation import FitChangeImplantLocation

@@ -18,7 +19,10 @@ class GuiAddImplantCommand(wx.Command):
        self.itemID = itemID

    def Do(self):
-        if self.internal_history.Submit(FitAddImplantCommand(self.fitID, self.itemID, True)) and self.internal_history.Submit(FitChangeImplantLocation(self.fitID, ImplantLocation.FIT)):
+        if (
+            self.internal_history.Submit(FitAddImplantCommand(fitID=self.fitID, implantInfo=ImplantInfo(itemID=self.itemID))) and
+            self.internal_history.Submit(FitChangeImplantLocation(self.fitID, ImplantLocation.FIT))
+        ):
            self.sFit.recalc(self.fitID)
            wx.PostEvent(self.mainFrame, GE.FitChanged(fitID=self.fitID))
            return True
--- a/gui/fitCommands/helpers.py
+++ b/gui/fitCommands/helpers.py
@@ -5,6 +5,7 @@ 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.implant import Implant
 from eos.saveddata.fighter import Fighter
 from eos.saveddata.module import Module
 from service.market import Market
@@ -90,53 +91,6 @@ class ModuleInfo:
            'chargeID', 'state', 'spoolType', 'spoolAmount'])


-class BoosterInfo:
-
-    def __init__(self, itemID, state=None, sideEffects=None):
-        self.itemID = itemID
-        self.state = state
-        self.sideEffects = sideEffects
-
-    @classmethod
-    def fromBooster(cls, booster):
-        if booster is None:
-            return None
-        info = cls(
-            itemID=booster.itemID,
-            state=booster.active,
-            sideEffects={se.effectID: se.active for se in booster.sideEffects})
-        return info
-
-    def toBooster(self):
-        item = Market.getInstance().getItem(self.itemID)
-        try:
-            booster = Booster(item)
-        except ValueError:
-            pyfalog.warning('Invalid item: {}'.format(self.itemID))
-            return None
-        if self.state is not None:
-            booster.active = self.state
-        if self.sideEffects is not None:
-            for sideEffect in booster.sideEffects:
-                sideEffect.active = self.sideEffects.get(sideEffect.effectID, sideEffect.active)
-        return booster
-
-class CargoInfo:
-
-    def __init__(self, itemID, amount):
-        self.itemID = itemID
-        self.amount = amount
-
-    def toCargo(self):
-        item = Market.getInstance().getItem(self.itemID)
-        cargo = Cargo(item)
-        cargo.amount = self.amount
-        return cargo
-
-    def __repr__(self):
-        return makeReprStr(self, ['itemID', 'amount'])
-
-
 class DroneInfo:

    def __init__(self, itemID, amount, amountActive):
@@ -208,6 +162,81 @@ class FighterInfo:
        return makeReprStr(self, ['itemID', 'amount', 'state', 'abilities'])


+class BoosterInfo:
+
+    def __init__(self, itemID, state=None, sideEffects=None):
+        self.itemID = itemID
+        self.state = state
+        self.sideEffects = sideEffects
+
+    @classmethod
+    def fromBooster(cls, booster):
+        if booster is None:
+            return None
+        info = cls(
+            itemID=booster.itemID,
+            state=booster.active,
+            sideEffects={se.effectID: se.active for se in booster.sideEffects})
+        return info
+
+    def toBooster(self):
+        item = Market.getInstance().getItem(self.itemID)
+        try:
+            booster = Booster(item)
+        except ValueError:
+            pyfalog.warning('Invalid item: {}'.format(self.itemID))
+            return None
+        if self.state is not None:
+            booster.active = self.state
+        if self.sideEffects is not None:
+            for sideEffect in booster.sideEffects:
+                sideEffect.active = self.sideEffects.get(sideEffect.effectID, sideEffect.active)
+        return booster
+
+
+class ImplantInfo:
+
+    def __init__(self, itemID, state=None):
+        self.itemID = itemID
+        self.state = state
+
+    @classmethod
+    def fromImplant(cls, implant):
+        if implant is None:
+            return None
+        info = cls(
+            itemID=implant.itemID,
+            state=implant.active)
+        return info
+
+    def toImplant(self):
+        item = Market.getInstance().getItem(self.itemID)
+        try:
+            implant = Implant(item)
+        except ValueError:
+            pyfalog.warning('Invalid item: {}'.format(self.itemID))
+            return None
+        if self.state is not None:
+            implant.active = self.state
+        return implant
+
+
+class CargoInfo:
+
+    def __init__(self, itemID, amount):
+        self.itemID = itemID
+        self.amount = amount
+
+    def toCargo(self):
+        item = Market.getInstance().getItem(self.itemID)
+        cargo = Cargo(item)
+        cargo.amount = self.amount
+        return cargo
+
+    def __repr__(self):
+        return makeReprStr(self, ['itemID', 'amount'])
+
+
 def stateLimit(itemIdentity):
    item = Market.getInstance().getItem(itemIdentity)
    if {'moduleBonusAssaultDamageControl', 'moduleBonusIndustrialInvulnerability'}.intersection(item.effects):