diff --git a/CHANGELOG.md b/CHANGELOG.md index d5982b1..4e97a82 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/src/budge/account.py b/src/budge/account.py index f086338..22c89ac 100644 --- a/src/budge/account.py +++ b/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], + ) + ) diff --git a/src/budge/date.py b/src/budge/date.py new file mode 100644 index 0000000..c8a6b2e --- /dev/null +++ b/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) + ) diff --git a/tests/budge/test_account.py b/tests/budge/test_account.py index 86f9391..7a6cfbb 100644 --- a/tests/budge/test_account.py +++ b/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)) diff --git a/tests/budge/test_date.py b/tests/budge/test_date.py new file mode 100644 index 0000000..6b11212 --- /dev/null +++ b/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), + ]