Add Account.balance_iter

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - Renamed `RecurringTransaction` and `RecurringTransfer` to
   `RepeatingTransaction` and `RepeatingTransfer`.
+- Add `Account.balance_iter` to iterate through balances by date
 
 ## v0.3.0 (2024-12-08)
 

src/budge/account.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from datetime import date
 from heapq import merge
+from itertools import groupby
 
 from stockholm import Money
 
@@ -24,11 +25,7 @@ def __iter__(self):
         by repeating transactions, ordered by date. This is useful for
         calculating or forecasting a balance for any point in time.
         """
-
-        for transaction in merge(
-            *self.repeating_transactions, sorted(self.transactions)
-        ):
-            yield transaction
+        yield from merge(sorted(self.transactions), *self.repeating_transactions)
 
     def transactions_range(
         self, start_date: date | None = None, end_date: date | None = None
@@ -49,3 +46,28 @@ def balance(self, as_of: date | None = None):
             transaction.amount
             for transaction in self.transactions_range(end_date=as_of)
         )
+
+    def balance_iter(
+        self, start_date: date | None = None, end_date: date | None = None
+    ):
+        """
+        Iterate over the account's balance over the given range, yielding a
+        tuple of each date in the range and the account balance on that date.
+
+        If `start_date` is not provided, the first yield will be the initial
+        balance of the account.
+
+        If `end_date` is not provided, the iteration will continue until all
+        transactions in the account have been iterated over.
+
+        :param start_date: The start date of the range
+        :param end_date: The end date of the range
+        :yield: A tuple of (date, balance) for each day in the range
+        """
+        bal = self.balance(start_date) if start_date else Money(0)
+
+        for date_, transactions in groupby(
+            self.transactions_range(start_date, end_date), lambda t: t.date
+        ):
+            bal += Money.sum(transaction.amount for transaction in transactions)
+            yield date_, bal

tests/budge/test_account.py

@@ -45,3 +45,22 @@ def test_transactions_range(self):
 
         transactions = list(self.acct.transactions_range(start_date, end_date))
         assert len(transactions) == 6
+
+    def test_balance_iter(self):
+        """
+        Verify that the balance_iter method returns the correct number of
+        balances between the given start and end dates.
+        """
+        start_date = self.today + relativedelta(months=6)
+        end_date = self.today + relativedelta(months=9)
+
+        balances = list(self.acct.balance_iter(start_date, end_date))
+
+        assert balances == [
+            (date(2023, 6, 15), Money(21)),
+            (date(2023, 7, 1), Money(22)),
+            (date(2023, 7, 15), Money(24)),
+            (date(2023, 8, 1), Money(25)),
+            (date(2023, 8, 15), Money(27)),
+            (date(2023, 9, 1), Money(28)),
+        ]