decimal-or-float, conversion

1 files changed, 6 insertions, 4 deletions

diff --git a/trade.py b/trade.py
index eedc9c4..e46d2b6 100644
--- a/trade.py
+++ b/trade.py
@@ -5,7 +5,7 @@ class Trade:
     Represents a cryptocurrency trade, including the amount traded, total cost, and the date of trade.
     Provides methods to modify the trade and access various attributes.
     """
-    def __init__(self, amount: Decimal, total_cost: Decimal, date: str) -> None:
+    def __init__(self, amount: float|Decimal, total_cost: float|Decimal, date: str) -> None:
         """
         Initialize a new Trade instance.
 
@@ -18,11 +18,11 @@ class Trade:
         if amount <= 0 or total_cost <= 0:
             raise ValueError("Amount and total cost must be > 0")
 
-        self.__amount: Decimal = amount
-        self.__total_cost: Decimal = total_cost
+        self.__amount: Decimal = Decimal(amount)
+        self.__total_cost: Decimal = Decimal(total_cost)
         self.__date: str = date
 
-    def remove_coins(self, amount: Decimal) -> None:
+    def remove_coins(self, amount: float|Decimal) -> None:
         """
         Reduce the amount of cryptocurrency in the trade by a specified amount.
 
@@ -35,6 +35,7 @@ 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
 
@@ -81,6 +82,7 @@ class Trade:
         """
         if self.amount == 0:
             raise ZeroDivisionError("Price per coin cannot be calculated when the amount is zero")
+
         return self.total_cost / self.amount
 
     def __repr__(self) -> str: