data: Add balance_of() function.

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):

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 <https://www.gnu.org/licenses/>.
+
+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

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,