Cache queue total

1 files changed, 12 insertions, 2 deletions

diff --git a/trade_queue.py b/trade_queue.py
index 0252431..08e0d06 100644
--- a/trade_queue.py
+++ b/trade_queue.py
@@ -13,6 +13,8 @@ class FIFOQueue:
     """
     def __init__(self) -> None:
         self.queue: Deque[Trade] = deque()
+        self._cached_total: Decimal = Decimal(0)
+        self._cache_valid: bool = True
 
     def add(self, amount: float|Decimal, total_cost: float|Decimal, date: str) -> None:
         """
@@ -20,6 +22,7 @@ class FIFOQueue:
         """
         trade = Trade(amount, total_cost, date)
         self.queue.append(trade)
+        self._cache_valid = False
 
     def remove_coins(self, amount: float|Decimal) -> List[Trade]:
         """
@@ -34,6 +37,8 @@ class FIFOQueue:
         if amount > self.get_remaining_amount():
             raise ValueError(f"Insufficient assets in queue to process sale of {amount}.")
 
+        self._cache_valid = False
+
         remaining: Decimal = amount
         entries: List[Trade] = []
 
@@ -44,7 +49,7 @@ class FIFOQueue:
                 ppc = trade.price_per_coin
                 trade.remove_coins(remaining)
                 entries.append(Trade(remaining, remaining * ppc, trade.date))
-                remaining = Decimal(0)
+                break
             else:
                 remaining -= trade.amount
                 entries.append(trade)
@@ -56,4 +61,9 @@ class FIFOQueue:
         """
         Calculate the total remaining amount in the queue.
         """
-        return sum((trade.amount for trade in self.queue), Decimal(0))
+
+        if not self._cache_valid:
+            self._cached_total = sum((trade.amount for trade in self.queue), Decimal(0))
+            self._cache_valid = True
+
+        return self._cached_total