data: Make balance_of currency-aware.

This commit is contained in:
Brett Smith 2020-04-09 14:13:07 -04:00
parent c6dc2d83ac
commit d66ba8773f
4 changed files with 48 additions and 19 deletions

View file

@ -25,6 +25,7 @@ import operator
from beancount.core import account as bc_account from beancount.core import account as bc_account
from beancount.core import amount as bc_amount from beancount.core import amount as bc_amount
from beancount.core import convert as bc_convert
from typing import ( from typing import (
cast, cast,
@ -257,18 +258,28 @@ class Posting(BasePosting):
def balance_of(txn: Transaction, def balance_of(txn: Transaction,
*preds: Callable[[Account], Optional[bool]], *preds: Callable[[Account], Optional[bool]],
) -> decimal.Decimal: ) -> Amount:
"""Return the balance of specified postings in a transaction. """Return the balance of specified postings in a transaction.
Given a transaction and a series of account predicates, balance_of Given a transaction and a series of account predicates, balance_of
returns the balance of the amounts of all postings with accounts that returns the balance of the amounts of all postings with accounts that
match any of the predicates. match any of the predicates.
balance_of uses the "weight" of each posting, so the return value will
use the currency of the postings' cost when available.
""" """
return sum( match_posts = [post for post in iter_postings(txn)
(post.units.number for post in iter_postings(txn) if any(pred(post.account) for pred in preds)]
if any(pred(post.account) for pred in preds)), number = decimal.Decimal(0)
decimal.Decimal(0), if not match_posts:
) currency = ''
else:
weights: Sequence[Amount] = [
bc_convert.get_weight(post) for post in match_posts # type:ignore[no-untyped-call]
]
number = sum((wt.number for wt in weights), number)
currency = weights[0].currency
return Amount._make((number, currency))
def iter_postings(txn: Transaction) -> Iterator[Posting]: def iter_postings(txn: Transaction) -> Iterator[Posting]:
"""Yield an enhanced Posting object for every posting in the transaction""" """Yield an enhanced Posting object for every posting in the transaction"""

View file

@ -40,7 +40,7 @@ class MetaApproval(core._RequireLinksPostingMetadataHook):
txn, txn,
data.Account.is_cash_equivalent, data.Account.is_cash_equivalent,
data.Account.is_credit_card, data.Account.is_credit_card,
) ).number
) )
def _run_on_post(self, txn: Transaction, post: data.Posting) -> bool: def _run_on_post(self, txn: Transaction, post: data.Posting) -> bool:

View file

@ -24,6 +24,7 @@ from . import testutil
from conservancy_beancount import data from conservancy_beancount import data
is_cash_eq = data.Account.is_cash_equivalent is_cash_eq = data.Account.is_cash_equivalent
USD = testutil.Amount
@pytest.fixture @pytest.fixture
def payable_payment_txn(): def payable_payment_txn():
@ -34,6 +35,14 @@ def payable_payment_txn():
('Assets:Checking', -5), ('Assets:Checking', -5),
]) ])
@pytest.fixture
def fx_donation_txn():
return testutil.Transaction(postings=[
('Income:Donations', -500, 'EUR', ('.9', 'USD')),
('Assets:Checking', 445),
('Expenses:BankingFees', 5),
])
def balance_under(txn, *accts): def balance_under(txn, *accts):
pred = methodcaller('is_under', *accts) pred = methodcaller('is_under', *accts)
return data.balance_of(txn, pred) return data.balance_of(txn, pred)
@ -43,17 +52,18 @@ def test_balance_of_simple_txn():
('Assets:Cash', 50), ('Assets:Cash', 50),
('Income:Donations', -50), ('Income:Donations', -50),
]) ])
assert balance_under(txn, 'Assets') == 50 assert balance_under(txn, 'Assets') == USD(50)
assert balance_under(txn, 'Income') == -50 assert balance_under(txn, 'Income') == USD(-50)
def test_zero_balance_of(payable_payment_txn): def test_zero_balance_of(payable_payment_txn):
assert balance_under(payable_payment_txn, 'Equity') == 0 expected = testutil.Amount(0, '')
assert balance_under(payable_payment_txn, 'Assets:Cash') == 0 assert balance_under(payable_payment_txn, 'Equity') == expected
assert balance_under(payable_payment_txn, 'Liabilities:CreditCard') == 0 assert balance_under(payable_payment_txn, 'Assets:Cash') == expected
assert balance_under(payable_payment_txn, 'Liabilities:CreditCard') == expected
def test_nonzero_balance_of(payable_payment_txn): def test_nonzero_balance_of(payable_payment_txn):
assert balance_under(payable_payment_txn, 'Assets', 'Expenses') == -50 assert balance_under(payable_payment_txn, 'Assets', 'Expenses') == USD(-50)
assert balance_under(payable_payment_txn, 'Assets', 'Liabilities') == -5 assert balance_under(payable_payment_txn, 'Assets', 'Liabilities') == USD(-5)
def test_multiarg_balance_of(): def test_multiarg_balance_of():
txn = testutil.Transaction(postings=[ txn = testutil.Transaction(postings=[
@ -61,7 +71,12 @@ def test_multiarg_balance_of():
('Expenses:BankingFees', 5), ('Expenses:BankingFees', 5),
('Assets:Checking', -655), ('Assets:Checking', -655),
]) ])
assert data.balance_of(txn, is_cash_eq, data.Account.is_credit_card) == -5 assert data.balance_of(txn, is_cash_eq, data.Account.is_credit_card) == USD(-5)
def test_balance_of_multipost_txn(payable_payment_txn): def test_balance_of_multipost_txn(payable_payment_txn):
assert data.balance_of(payable_payment_txn, is_cash_eq) == -55 assert data.balance_of(payable_payment_txn, is_cash_eq) == USD(-55)
def test_balance_of_multicurrency_txn(fx_donation_txn):
assert balance_under(fx_donation_txn, 'Income') == USD(-450)
assert balance_under(fx_donation_txn, 'Income', 'Assets') == USD(-5)
assert balance_under(fx_donation_txn, 'Income', 'Expenses') == USD(-445)

View file

@ -69,16 +69,19 @@ def test_path(s):
def Amount(number, currency='USD'): def Amount(number, currency='USD'):
return bc_amount.Amount(Decimal(number), currency) return bc_amount.Amount(Decimal(number), currency)
def Cost(number, currency='USD', date=FY_MID_DATE, label=None):
return bc_data.Cost(Decimal(number), currency, date, label)
def Posting(account, number, def Posting(account, number,
currency='USD', cost=None, price=None, flag=None, currency='USD', cost=None, price=None, flag=None,
**meta): **meta):
if not (number is None or isinstance(number, Decimal)): if cost is not None:
number = Decimal(number) cost = Cost(*cost)
if meta is None: if meta is None:
meta = None meta = None
return bc_data.Posting( return bc_data.Posting(
account, account,
bc_amount.Amount(number, currency), Amount(number, currency),
cost, cost,
price, price,
flag, flag,