Change Account.balance_iter to Account.daily_balance

CHANGELOG.md

@@ -2,9 +2,9 @@
 
 ## v0.4.0 (unreleased)
 
+- Added `Account.daily_balance` to iterate through account balance by date.
 - 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,10 +1,12 @@
 from dataclasses import dataclass, field
-from datetime import date
+from datetime import date, timedelta
 from heapq import merge
 from itertools import groupby
+from typing import Generator
 
 from stockholm import Money
 
+from .date import daterange
 from .transaction import RepeatingTransaction, Transaction
 
 
@@ -47,27 +49,61 @@ def balance(self, as_of: date | None = None):
             for transaction in self.transactions_range(end_date=as_of)
         )
 
-    def balance_iter(
+    def daily_balance(
         self, start_date: date | None = None, end_date: date | None = None
-    ):
+    ) -> Generator[tuple[date, Money]]:
+        """
+        Iterate over the daily balance of the account, yielding tuples of date
+        and balance.
+
+        The balance on the given start date is yielded first, and then the
+        balance for each subsequent date is yielded. If the start date is not
+        given, the date of the first transaction is used. If the end date is not
+        given, today's date is used.
         """
-        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.
+        start_date = start_date or next(self.transactions_range()).date
+        end_date = end_date or date.today()
 
-        If `start_date` is not provided, the first yield will be the initial
-        balance of the account.
+        balance = self.balance(start_date)
+        yield start_date, balance
 
-        If `end_date` is not provided, the iteration will continue until all
-        transactions in the account have been iterated over.
+        for _date, delta in self._daily_balance_delta(
+            start_date + timedelta(days=1), end_date
+        ):
+            balance += delta
+            yield _date, balance
 
-        :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
+    def _deltas_by_date(self, start_date: date, end_date: date):
         """
-        bal = self.balance(start_date) if start_date else Money(0)
+        Iterate over the deltas in the account balance for each date in the
+        given range, including the given start and end dates.
 
-        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
+        Yields tuples, where the first element is the date and the second
+        element is the total amount of all transactions on that date.
+        """
+        return (
+            (_date, Money.sum(transaction.amount for transaction in transactions))
+            for _date, transactions in groupby(
+                self.transactions_range(start_date, end_date), key=lambda t: t.date
+            )
+        )
+
+    def _daily_balance_delta(self, start_date: date, end_date: date):
+        """
+        Calculate the daily change in account balance over the specified date range.
+
+        This function yields tuples, where the first element is the date and the
+        second element is the net change in balance for that date. It combines the
+        deltas from actual transactions and placeholder deltas for dates without
+        transactions to ensure a continuous range.
+        """
+        return (
+            (_date, Money.sum(delta[1] for delta in deltas))
+            for _date, deltas in groupby(
+                merge(
+                    self._deltas_by_date(start_date, end_date),
+                    ((_date, Money(0)) for _date in daterange(start_date, end_date)),
+                ),
+                key=lambda t: t[0],
+            )
+        )

src/budge/date.py

@@ -0,0 +1,8 @@
+from datetime import date, timedelta
+
+
+def daterange(start_date: date, end_date: date):
+    """Iterate over a range of dates, inclusive of the start and end dates."""
+    return (
+        start_date + timedelta(days=x) for x in range((end_date - start_date).days + 1)
+    )

tests/budge/test_account.py

@@ -10,7 +10,7 @@
 class TestAccount:
     today = date(2022, 12, 6)
 
-    t1 = Transaction(Money(1), "test 1", date(2022, 12, 6))
+    t1 = Transaction(Money(1), "test 1", date(2022, 12, 1))
 
     rule1 = rrule(freq=MONTHLY, bymonthday=1, dtstart=today)
     rt1 = RepeatingTransaction(Money(1), "test 1", schedule=rule1)
@@ -46,21 +46,35 @@ def test_transactions_range(self):
         transactions = list(self.acct.transactions_range(start_date, end_date))
         assert len(transactions) == 6
 
-    def test_balance_iter(self):
+    def test_daily_balance_past(self):
         """
-        Verify that the balance_iter method returns the correct number of
-        balances between the given start and end dates.
+        Verify that the daily_balance method returns the correct balances for each
+        day in the past month, starting from a given start date and ending on today's
+        date. The initial balance should be zero, and the balance on today's date
+        should match the expected value.
         """
-        start_date = self.today + relativedelta(months=6)
-        end_date = self.today + relativedelta(months=9)
+        start_date = date(2022, 11, 6)
+        balances = list(
+            self.acct.daily_balance(start_date=start_date, end_date=self.today)
+        )
+
+        assert len(balances) == 31
+        assert balances[0] == (start_date, Money(0))
+        assert balances[-1] == (self.today, Money(1))
 
-        balances = list(self.acct.balance_iter(start_date, end_date))
+    def test_daily_balance_future(self):
+        """
+        Verify that the daily_balance method returns the correct balances for each
+        day in the future month, starting from today's date and ending on a given
+        end date. The initial balance should be the expected value, and the balance
+        on the end date should match the expected value.
+        """
+        end_date = self.today + relativedelta(months=1)
+        balances = list(
+            self.acct.daily_balance(start_date=self.today, end_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)),
-        ]
+        assert len(balances) == 32
+        assert balances[0] == (self.today, Money(1))
+        assert balances[9] == (date(2022, 12, 15), Money(3))
+        assert balances[-1] == (end_date, Money(4))

tests/budge/test_date.py

@@ -0,0 +1,23 @@
+from datetime import date
+
+from dateutil.relativedelta import relativedelta
+
+from budge.date import daterange
+
+
+def test_daterange():
+    start_date = date(2022, 12, 28)
+    end_date = start_date + relativedelta(weeks=1)
+
+    dates = list(daterange(start_date, end_date))
+
+    assert dates == [
+        date(2022, 12, 28),
+        date(2022, 12, 29),
+        date(2022, 12, 30),
+        date(2022, 12, 31),
+        date(2023, 1, 1),
+        date(2023, 1, 2),
+        date(2023, 1, 3),
+        date(2023, 1, 4),
+    ]