From c01307202a8125a222b2600b4cbcaa746591164d Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <yohann.danello@gmail.com>
Date: Fri, 18 Dec 2020 20:01:52 +0100
Subject: [PATCH] Add shields to be more protected, see #48

---
 squirrelbattle/display/texturepack.py |  2 ++
 squirrelbattle/entities/items.py      | 50 ++++++++++++++++++++-------
 2 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/squirrelbattle/display/texturepack.py b/squirrelbattle/display/texturepack.py
index 1f6dc76..2295667 100644
--- a/squirrelbattle/display/texturepack.py
+++ b/squirrelbattle/display/texturepack.py
@@ -70,6 +70,7 @@ TexturePack.ASCII_PACK = TexturePack(
     MERCHANT='M',
     PLAYER='@',
     RABBIT='Y',
+    SHIELD='D',
     SUNFLOWER='I',
     SWORD='\u2020',
     TEDDY_BEAR='8',
@@ -96,6 +97,7 @@ TexturePack.SQUIRREL_PACK = TexturePack(
     PLAYER='🐿️ ️',
     MERCHANT='🦜',
     RABBIT='🐇',
+    SHIELD='🛡️ ',
     SUNFLOWER='🌻',
     SWORD='🗡️ ',
     TEDDY_BEAR='🧸',
diff --git a/squirrelbattle/entities/items.py b/squirrelbattle/entities/items.py
index 4da4a8c..b3f8a84 100644
--- a/squirrelbattle/entities/items.py
+++ b/squirrelbattle/entities/items.py
@@ -40,14 +40,20 @@ class Item(Entity):
         Indicates what should be done when the item is used.
         """
 
-    def equip(self) -> None:
+    def equip(self, armor: bool = False) -> None:
         """
         Indicates what should be done when the item is equipped.
         """
-        if self.held_by.equipped_item:
-            self.held_by.equipped_item.unequip()
-        self.held_by.remove_from_inventory(self)
-        self.held_by.equipped_item = self
+        if armor:
+            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
+        else:
+            if self.held_by.equipped_item:
+                self.held_by.equipped_item.unequip()
+            self.held_by.remove_from_inventory(self)
+            self.held_by.equipped_item = self
 
     def unequip(self) -> None:
         """
@@ -75,7 +81,7 @@ class Item(Entity):
 
     @staticmethod
     def get_all_items() -> list:
-        return [BodySnatchPotion, Bomb, Heart, Sword]
+        return [BodySnatchPotion, Bomb, Heart, Shield, Sword]
 
     def be_sold(self, buyer: InventoryHolder, seller: InventoryHolder) -> bool:
         """
@@ -225,20 +231,17 @@ class Sword(Weapon):
     """
     A basic weapon
     """
-    strength: int
-
-    def __init__(self, name: str = "sword", price: int = 20, strength: int = 3,
+    def __init__(self, name: str = "sword", price: int = 20,
                  *args, **kwargs):
         super().__init__(name=name, price=price, *args, **kwargs)
         self.name = name
-        self.strength = strength
 
-    def equip(self) -> None:
+    def equip(self, armor: bool = False) -> None:
         """
         When a sword is equipped, the player gains strength.
         """
         super().equip()
-        self.held_by.strength += self.strength
+        self.held_by.strength += self.damage
 
     def unequip(self) -> None:
         """
@@ -246,7 +249,28 @@ class Sword(Weapon):
         :return:
         """
         super().unequip()
-        self.held_by.strength -= self.strength
+        self.held_by.strength -= self.damage
+
+
+class Shield(Item):
+    constitution: int
+
+    def __init__(self, constitution: int = 2, *args, **kwargs):
+        super().__init__(name="shield", *args, **kwargs)
+        self.constitution = constitution
+
+    def equip(self, armor: bool = True) -> None:
+        super().equip(armor)
+        self.held_by.constitution += self.constitution
+
+    def unequip(self) -> None:
+        super().unequip()
+        self.held_by.constitution -= self.constitution
+
+    def save_state(self) -> dict:
+        d = super().save_state()
+        d["constitution"] = self.constitution
+        return d
 
 
 class BodySnatchPotion(Item):