diff --git a/conservancy_beancount/data.py b/conservancy_beancount/data.py index 8e191dfbf6677cc57602ab4f1d21d09fdc621fdf..b16e60e8856a9fdf0cc18ea77a72d4178080eb1d 100644 --- a/conservancy_beancount/data.py +++ b/conservancy_beancount/data.py @@ -260,6 +260,31 @@ class Posting(BasePosting): return self.account.is_cash_equivalent() and self.is_debit(threshold, default) +def balance_of(txn: Transaction, + *accounts: str, + default: Optional[DecimalCompat]=None, +) -> Optional[decimal.Decimal]: + """Return the balance of specified postings in a transaction. + + Given a transaction and a series of account names, balance_of returns the + balance of the amounts of all postings under those account names. + + Account names are matched using Account.is_under. Refer to that docstring + for details about what matches. + + If any of the postings have no amount, returns default. + """ + if default is not None: + default = decimal.Decimal(default) + retval = decimal.Decimal(0) + for post in txn.postings: + if Account(post.account).is_under(*accounts): + if post.units.number is None: + return default + else: + retval += post.units.number + return retval + def iter_postings(txn: Transaction) -> Iterator[Posting]: """Yield an enhanced Posting object for every posting in the transaction""" for index, source in enumerate(txn.postings): diff --git a/tests/test_data_balance_of.py b/tests/test_data_balance_of.py new file mode 100644 index 0000000000000000000000000000000000000000..3ec251a93d6431b9b5d483a7c8215425f38b06e4 --- /dev/null +++ b/tests/test_data_balance_of.py @@ -0,0 +1,90 @@ +"""Test data.balance_of function""" +# Copyright © 2020 Brett Smith +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +from decimal import Decimal + +import pytest + +from . import testutil + +from conservancy_beancount import data + +@pytest.fixture +def payable_payment_txn(): + return testutil.Transaction(postings=[ + ('Liabilities:Payable:Accounts', 50), + ('Assets:Checking', -50), + ('Expenses:BankingFees', 5), + ('Assets:Checking', -5), + ]) + +@pytest.fixture +def none_posting_txn(): + return testutil.Transaction(postings=[ + ('Income:Donations', -30), + ('Expenses:BankingFees', 3), + ('Assets:Checking', None), + ]) + +@pytest.fixture +def multipost_one_none_txn(): + return testutil.Transaction(postings=[ + ('Liabilities:Payable:Accounts', 50), + ('Assets:Checking', -50), + ('Expenses:BankingFees', 5), + ('Assets:Checking', None), + ]) + +def test_balance_of_simple_txn(): + txn = testutil.Transaction(postings=[ + ('Assets:Cash', 50), + ('Income:Donations', -50), + ]) + assert data.balance_of(txn, 'Assets') == 50 + assert data.balance_of(txn, 'Income') == -50 + +def test_zero_balance_of(payable_payment_txn): + assert data.balance_of(payable_payment_txn, 'Equity') == 0 + assert data.balance_of(payable_payment_txn, 'Assets:Cash') == 0 + assert data.balance_of(payable_payment_txn, 'Liabilities:CreditCard') == 0 + +def test_balance_of_multipost_txn(payable_payment_txn): + assert data.balance_of(payable_payment_txn, 'Assets') == -55 + +def test_multiarg_balance_of(payable_payment_txn): + assert data.balance_of(payable_payment_txn, 'Assets', 'Expenses') == -50 + assert data.balance_of(payable_payment_txn, 'Assets', 'Liabilities') == -5 + +def test_balance_of_uses_whole_account_names(payable_payment_txn): + assert data.balance_of(payable_payment_txn, 'Assets:Check') == 0 + +def test_balance_of_none_posting(none_posting_txn): + assert data.balance_of(none_posting_txn, 'Assets') is None + +def test_balance_of_none_posting_with_default(none_posting_txn): + expected = Decimal('Infinity') + assert data.balance_of(none_posting_txn, 'Assets', default=expected) == expected + +def test_balance_of_other_side_of_none_posting(none_posting_txn): + assert data.balance_of(none_posting_txn, 'Income') == -30 + assert data.balance_of(none_posting_txn, 'Expenses') == 3 + +def test_balance_of_multi_postings_one_none(multipost_one_none_txn): + assert data.balance_of(multipost_one_none_txn, 'Assets') is None + +def test_balance_of_multi_postings_one_none(multipost_one_none_txn): + expected = Decimal('Infinity') + assert data.balance_of(multipost_one_none_txn, 'Assets', default=expected) == expected diff --git a/tests/testutil.py b/tests/testutil.py index cb6fc1b384936b5a21b0f8cf93a12b1121935e15..583de75a1a9cd2ed47bc37cad028c499c2984e76 100644 --- a/tests/testutil.py +++ b/tests/testutil.py @@ -72,11 +72,13 @@ def Amount(number, currency='USD'): def Posting(account, number, currency='USD', cost=None, price=None, flag=None, **meta): - if not meta: + if not (number is None or isinstance(number, Decimal)): + number = Decimal(number) + if meta is None: meta = None return bc_data.Posting( account, - bc_amount.Amount(Decimal(number), currency), + bc_amount.Amount(number, currency), cost, price, flag,