Fix EquipType in case of Weapons

diff --git a/MyCollection/Internals/Classes/Item.lua b/MyCollection/Internals/Classes/Item.lua
index eb14758..87e05b4 100644
--- a/MyCollection/Internals/Classes/Item.lua
+++ b/MyCollection/Internals/Classes/Item.lua
@@ -100,6 +100,7 @@ function Item:Initialize(characterId, bagId, slotId, setId, equipType, traitType
         else
             self.weaponType = weaponType
         end
+        self.equipType = Extensions.Constants.GetEquipTypeOfWeapon(self.weaponType)
     else
         if Extensions.Constants.IsArmor(self.equipType) then
             if armorType == nil then
diff --git a/MyCollection/Internals/Classes/Piece.lua b/MyCollection/Internals/Classes/Piece.lua
index ab2a040..ee16bc6 100644
--- a/MyCollection/Internals/Classes/Piece.lua
+++ b/MyCollection/Internals/Classes/Piece.lua
@@ -1,5 +1,6 @@
 local Classes = MyCollection.Internals.Classes
 local Constants = MyCollection.Internals.Constants
+local Extensions = MyCollection.Internals.Functions.Extensions
 local Dependencies = MyCollection.Internals.Dependencies
 local Logger = MyCollection.Internals.Dependencies.Logger
 local Data = MyCollection.Internals.Data
@@ -79,6 +80,10 @@ function Piece:Initialize(setId, equipType, traitType, armorType, weaponType)
     self.armorType = armorType
     self.weaponType = weaponType

+    if self.weaponType ~= nil then
+        self.equipType = Extensions.Constants.GetEquipTypeOfWeapon(self.weaponType)
+    end
+
     self.references = {}

     self:Refresh()
diff --git a/MyCollection/Internals/Functions/Extensions/Constants.lua b/MyCollection/Internals/Functions/Extensions/Constants.lua
index 17f2727..8883c0b 100644
--- a/MyCollection/Internals/Functions/Extensions/Constants.lua
+++ b/MyCollection/Internals/Functions/Extensions/Constants.lua
@@ -19,6 +19,22 @@ Extensions.Constants.IsArmor = function (equipType)
     return false
 end

+Extensions.Constants.GetEquipTypeOfWeapon = function (weaponType)
+    if weaponType == Constants.WeaponTypes.Axe then                 return Constants.EquipTypes.Weapons.OneHand end
+    if weaponType == Constants.WeaponTypes.Hammer then              return Constants.EquipTypes.Weapons.OneHand end
+    if weaponType == Constants.WeaponTypes.Sword then               return Constants.EquipTypes.Weapons.OneHand end
+    if weaponType == Constants.WeaponTypes.TwoHandedSword then      return Constants.EquipTypes.Weapons.TwoHand end
+    if weaponType == Constants.WeaponTypes.TwoHandedAxe then        return Constants.EquipTypes.Weapons.TwoHand end
+    if weaponType == Constants.WeaponTypes.TwoHandedHammer then     return Constants.EquipTypes.Weapons.TwoHand end
+    if weaponType == Constants.WeaponTypes.Bow then                 return Constants.EquipTypes.Weapons.TwoHand end
+    if weaponType == Constants.WeaponTypes.HealingStaff then        return Constants.EquipTypes.Weapons.TwoHand end
+    if weaponType == Constants.WeaponTypes.Dagger then              return Constants.EquipTypes.Weapons.OneHand end
+    if weaponType == Constants.WeaponTypes.FireStaff then           return Constants.EquipTypes.Weapons.TwoHand end
+    if weaponType == Constants.WeaponTypes.FrostStaff then          return Constants.EquipTypes.Weapons.TwoHand end
+    if weaponType == Constants.WeaponTypes.LightningStaff then      return Constants.EquipTypes.Weapons.TwoHand end
+    if weaponType == Constants.WeaponTypes.Shield then              return Constants.EquipTypes.Weapons.OffHand end
+end
+
 Extensions.Constants.IsWeaponOrShield = function (equipType)
     for key, id in pairs(Constants.EquipTypes.Weapons) do
         if (id == equipType) then