data.balance_of: Take account predicates, not just names.

For increased flexibility.
In particular, now you can pass in Account boolean methods to
call those directly.

conservancy_beancount/data.py

@@ -261,26 +261,23 @@ class Posting(BasePosting):
 
 
 def balance_of(txn: Transaction,
-               *accounts: str,
+               *preds: Callable[[Account], Optional[bool]],
                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.
+    Given a transaction and a series of account predicates, balance_of
+    returns the balance of the amounts of all postings with accounts that
+    match any of the predicates.
 
     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):
+        acct = Account(post.account)
+        if any(p(acct) for p in preds):
             if post.units.number is None:
-                return default
+                return None if default is None else decimal.Decimal(default)
             else:
                 retval += post.units.number
     return retval

tests/test_data_balance_of.py

@@ -15,6 +15,7 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 from decimal import Decimal
+from operator import methodcaller
 
 import pytest
 
@@ -22,6 +23,8 @@ from . import testutil
 
 from conservancy_beancount import data
 
+is_cash_eq = data.Account.is_cash_equivalent
+
 @pytest.fixture
 def payable_payment_txn():
     return testutil.Transaction(postings=[
@@ -48,43 +51,56 @@ def multipost_one_none_txn():
         ('Assets:Checking', None),
     ])
 
+def balance_under(txn, *accts):
+    pred = methodcaller('is_under', *accts)
+    return data.balance_of(txn, pred)
+
 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
+    assert balance_under(txn, 'Assets') == 50
+    assert balance_under(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
+    assert balance_under(payable_payment_txn, 'Equity') == 0
+    assert balance_under(payable_payment_txn, 'Assets:Cash') == 0
+    assert balance_under(payable_payment_txn, 'Liabilities:CreditCard') == 0
+
+def test_nonzero_balance_of(payable_payment_txn):
+    assert balance_under(payable_payment_txn, 'Assets', 'Expenses') == -50
+    assert balance_under(payable_payment_txn, 'Assets', 'Liabilities') == -5
+
+def test_multiarg_balance_of():
+    txn = testutil.Transaction(postings=[
+        ('Liabilities:CreditCard', 650),
+        ('Expenses:BankingFees', 5),
+        ('Assets:Checking', -655),
+    ])
+    assert data.balance_of(txn, is_cash_eq, data.Account.is_credit_card) == -5
 
 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
+    assert data.balance_of(payable_payment_txn, is_cash_eq) == -55
 
 def test_balance_of_none_posting(none_posting_txn):
-    assert data.balance_of(none_posting_txn, 'Assets') is None
+    assert data.balance_of(none_posting_txn, is_cash_eq) 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
+    assert expected == data.balance_of(
+        none_posting_txn, is_cash_eq, default=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
+    assert balance_under(none_posting_txn, 'Income') == -30
+    assert balance_under(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
+    assert data.balance_of(multipost_one_none_txn, is_cash_eq) 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
+    assert expected == data.balance_of(
+        multipost_one_none_txn, is_cash_eq, default=expected,
+    )