Add removal methods

1 files changed, 17 insertions, 5 deletions

diff --git a/trade.py b/trade.py
index 1a39b2d..ed646d1 100644
--- a/trade.py
+++ b/trade.py
@@ -1,4 +1,10 @@
 from decimal import Decimal
+from enum import Enum
+
+class PriceAdaption(Enum):
+    KeepTotalCost = 1
+    KeepPricePerCoin = 2
+
 
 class Trade:
     """
@@ -22,11 +28,11 @@ class Trade:
         self.__total_cost: Decimal = Decimal(total_cost)
         self.__date: str = date
 
-    def remove_coins(self, amount: float|Decimal) -> None:
+    def remove_coins(self, amount: float|Decimal, adapt: PriceAdaption = PriceAdaption.KeepPricePerCoin) -> None:
         """
         Reduce the amount of cryptocurrency in the trade by a specified amount.
 
-        This effectively "loses" coins. The price-per-coin remains the same.
+        This effectively "loses" coins.
 
         Args:
             amount (Decimal): The amount of cryptocurrency to remove.
@@ -37,9 +43,15 @@ class Trade:
         if amount > self.__amount:
             raise ValueError(f"Can't remove more than {self.__amount}")
         
-        amount = Decimal(amount)
-        self.__total_cost -= amount * self.price_per_coin
-        self.__amount -= amount
+        if adapt == PriceAdaption.KeepPricePerCoin:
+            amount = Decimal(amount)
+            self.__total_cost -= amount * self.price_per_coin
+            self.__amount -= amount
+        elif adapt == PriceAdaption.KeepTotalCost:
+            amount = Decimal(amount)
+            self.__amount -= amount
+        else:
+            raise ValueError("Unknown adaptation strategy")
 
     @property
     def amount(self) -> Decimal: