Limit fighters to fit's available slots. We hijack the module slot bit to do the dirty work >_>

eos/saveddata/fighter.py

@@ -21,11 +21,13 @@ from eos.modifiedAttributeDict import ModifiedAttributeDict, ItemAttrShortcut, C
 from eos.effectHandlerHelpers import HandledItem, HandledCharge, HandledDroneCargoList
 from sqlalchemy.orm import validates, reconstructor
 import eos.db
+from eos.enum import Enum
 import logging
-from eos.types import FighterAbility
+from eos.types import FighterAbility, Slot
 
 logger = logging.getLogger(__name__)
 
+
 class Fighter(HandledItem, HandledCharge, ItemAttrShortcut, ChargeAttrShortcut):
     DAMAGE_TYPES = ("em", "kinetic", "explosive", "thermal")
     DAMAGE_TYPES2 = ("EM", "Kin", "Exp", "Therm")
@@ -71,29 +73,45 @@ class Fighter(HandledItem, HandledCharge, ItemAttrShortcut, ChargeAttrShortcut):
         self.__dps = None
         self.__volley = None
         self.__miningyield = None
+        self.__itemModifiedAttributes = ModifiedAttributeDict()
+        self.__chargeModifiedAttributes = ModifiedAttributeDict()
 
         if len(self.abilities) != len(self.item.effects):
             self.__abilities = []
             for ability in self.__getAbilities():
                 self.__abilities.append(ability)
 
-        self.__itemModifiedAttributes = ModifiedAttributeDict()
-        self.__itemModifiedAttributes.original = self.__item.attributes
-        self.__itemModifiedAttributes.overrides = self.__item.overrides
+        if self.__item:
+            self.__itemModifiedAttributes.original = self.__item.attributes
+            self.__itemModifiedAttributes.overrides = self.__item.overrides
+            self.__slot = self.__calculateSlot(self.__item)
 
-        self.__chargeModifiedAttributes = ModifiedAttributeDict()
-        chargeID = self.getModifiedItemAttr("entityMissileTypeID")
-        if chargeID is not None:
-            charge = eos.db.getItem(int(chargeID))
-            self.__charge = charge
-            self.__chargeModifiedAttributes.original = charge.attributes
-            self.__chargeModifiedAttributes.overrides = charge.overrides
+            # chargeID = self.getModifiedItemAttr("fighterAbilityLaunchBombType")
+            chargeID = None
+            if chargeID is not None:
+                charge = eos.db.getItem(int(chargeID))
+                self.__charge = charge
+                self.__chargeModifiedAttributes.original = charge.attributes
+                self.__chargeModifiedAttributes.overrides = charge.overrides
 
     def __getAbilities(self):
         """Returns list of FighterAbilities that are loaded with data"""
         print "getting list of abilities"
         return [FighterAbility(effect) for effect in self.item.effects.values()]
 
+    def __calculateSlot(self, item):
+        types = {"Light": Slot.F_LIGHT,
+                 "Support": Slot.F_SUPPORT,
+                 "Heavy": Slot.F_HEAVY}
+
+        for t, slot in types.iteritems():
+            if self.getModifiedItemAttr("fighterSquadronIs{}".format(t)):
+                return slot
+
+    @property
+    def slot(self):
+        return self.__slot
+
     @property
     def amountActive(self):
         return self.getModifiedItemAttr("fighterSquadronMaxSize") if self.amount == -1 else self.amount
@@ -239,15 +257,8 @@ class Fighter(HandledItem, HandledCharge, ItemAttrShortcut, ChargeAttrShortcut):
         return copy
 
     def fits(self, fit):
+        if fit.getSlotsFree(self.slot) <= 0:
+            return False
+
         return True
 
-        fitDroneGroupLimits = set()
-        for i in xrange(1, 3):
-            groneGrp = fit.ship.getModifiedItemAttr("allowedDroneGroup%d" % i)
-            if groneGrp is not None:
-                fitDroneGroupLimits.add(int(groneGrp))
-        if len(fitDroneGroupLimits) == 0:
-            return True
-        if self.item.groupID in fitDroneGroupLimits:
-            return True
-        return False

eos/saveddata/fit.py

@@ -616,8 +616,8 @@ class Fit(object):
 
     def getSlotsUsed(self, type, countDummies=False):
         amount = 0
-        for mod in self.modules:
-            if mod.slot is type and (not mod.isEmpty or countDummies):
+        for mod in chain(self.modules, self.fighters):
+            if mod.slot is type and (not getattr(mod, "isEmpty", False) or countDummies):
                 amount += 1
 
         return amount
@@ -627,7 +627,10 @@ class Fit(object):
                  Slot.MED: "medSlots",
                  Slot.HIGH: "hiSlots",
                  Slot.RIG: "rigSlots",
-                 Slot.SUBSYSTEM: "maxSubSystems"}
+                 Slot.SUBSYSTEM: "maxSubSystems",
+                 Slot.F_LIGHT: "fighterLightSlots",
+                 Slot.F_SUPPORT: "fighterSupportSlots",
+                 Slot.F_HEAVY: "fighterHeavySlots"}
 
         if type in (Slot.MODE, Slot.SYSTEM):
             # These slots don't really exist, return default 0

eos/saveddata/module.py

@@ -46,6 +46,10 @@ class Slot(Enum):
     # system effects. They are projected "modules" and pyfa assumes all modules
     # have a slot. In this case, make one up.
     SYSTEM = 7
+    # fighter 'slots'. Just easier to put them here...
+    F_LIGHT = 10
+    F_SUPPORT = 11
+    F_HEAVY = 12
 
 class Hardpoint(Enum):
     NONE = 0