From 1a78ad584cd4c973a9de48d0bd87c374cc165b91 Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <yohann.danello@gmail.com>
Date: Wed, 6 Jan 2021 18:31:28 +0100
Subject: [PATCH] Move equip functions for items

---
 squirrelbattle/entities/items.py | 57 +++++++++++++++-----------------
 1 file changed, 27 insertions(+), 30 deletions(-)

diff --git a/squirrelbattle/entities/items.py b/squirrelbattle/entities/items.py
index 3eb3530..ee5f4a9 100644
--- a/squirrelbattle/entities/items.py
+++ b/squirrelbattle/entities/items.py
@@ -44,32 +44,11 @@ class Item(Entity):
         """
         Indicates what should be done when the item is equipped.
         """
-        if isinstance(self, Chestplate):
-            if self.held_by.equipped_armor:
-                self.held_by.equipped_armor.unequip()
-            self.held_by.remove_from_inventory(self)
-            self.held_by.equipped_armor = self
-        elif isinstance(self, Helmet):
-            if self.held_by.equipped_helmet:
-                self.held_by.equipped_helmet.unequip()
-            self.held_by.remove_from_inventory(self)
-            self.held_by.equipped_helmet = self
-        elif isinstance(self, Weapon):
-            if self.held_by.equipped_main:
-                if self.held_by.equipped_secondary:
-                    self.held_by.equipped_secondary.unequip()
-                self.held_by.remove_from_inventory(self)
-                self.held_by.equipped_secondary = self
-                # For weapons, they are equipped as main only if main is empty.
-            else:
-                self.held_by.remove_from_inventory(self)
-                self.held_by.equipped_main = self
-        else:
-            # Other objects are only equipped as secondary.
-            if self.held_by.equipped_secondary:
-                self.held_by.equipped_secondary.unequip()
-            self.held_by.remove_from_inventory(self)
-            self.held_by.equipped_secondary = self
+        # Other objects are only equipped as secondary.
+        if self.held_by.equipped_secondary:
+            self.held_by.equipped_secondary.unequip()
+        self.held_by.remove_from_inventory(self)
+        self.held_by.equipped_secondary = self
 
     def unequip(self) -> None:
         """
@@ -257,7 +236,8 @@ class Weapon(Item):
         """
         When a weapon is equipped, the player gains strength.
         """
-        super().equip()
+        self.held_by.remove_from_inventory(self)
+        self.held_by.equipped_main = self
         self.held_by.strength += self.damage
 
     def unequip(self) -> None:
@@ -309,7 +289,8 @@ class Shield(Armor):
 
     def __init__(self, name: str = "shield", constitution: int = 2,
                  price: int = 6, *args, **kwargs):
-        super().__init__(name=name, constitution=constitution, *args, **kwargs)
+        super().__init__(name=name, constitution=constitution, price=price,
+                         *args, **kwargs)
 
 
 class Helmet(Armor):
@@ -319,7 +300,15 @@ class Helmet(Armor):
 
     def __init__(self, name: str = "helmet", constitution: int = 2,
                  price: int = 8, *args, **kwargs):
-        super().__init__(name=name, constitution=constitution, *args, **kwargs)
+        super().__init__(name=name, constitution=constitution, price=price,
+                         *args, **kwargs)
+
+    def equip(self) -> None:
+        super().equip()
+        if self.held_by.equipped_helmet:
+            self.held_by.equipped_helmet.unequip()
+        self.held_by.remove_from_inventory(self)
+        self.held_by.equipped_helmet = self
 
 
 class Chestplate(Armor):
@@ -329,7 +318,15 @@ class Chestplate(Armor):
 
     def __init__(self, name: str = "chestplate", constitution: int = 4,
                  price: int = 15, *args, **kwargs):
-        super().__init__(name=name, constitution=constitution, *args, **kwargs)
+        super().__init__(name=name, constitution=constitution, price=price,
+                         *args, **kwargs)
+
+    def equip(self) -> None:
+        super().equip()
+        if self.held_by.equipped_armor:
+            self.held_by.equipped_armor.unequip()
+        self.held_by.remove_from_inventory(self)
+        self.held_by.equipped_armor = self
 
 
 class BodySnatchPotion(Item):