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 amount as bc_amount
from beancount.core import convert as bc_convert
from typing import (
cast,
@ -257,18 +258,28 @@ class Posting(BasePosting):
def balance_of(txn: Transaction,
*preds: Callable[[Account], Optional[bool]],
) -> decimal.Decimal:
) -> Amount:
"""Return the balance of specified postings in a transaction.
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.
balance_of uses the "weight" of each posting, so the return value will
use the currency of the postings' cost when available.
"""
return sum(
(post.units.number for post in iter_postings(txn)
if any(pred(post.account) for pred in preds)),
decimal.Decimal(0),
)
match_posts = [post for post in iter_postings(txn)
if any(pred(post.account) for pred in preds)]
number = 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]:
"""Yield an enhanced Posting object for every posting in the transaction"""

View file

@ -40,7 +40,7 @@ class MetaApproval(core._RequireLinksPostingMetadataHook):
txn,
data.Account.is_cash_equivalent,
data.Account.is_credit_card,
)
).number
)
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
is_cash_eq = data.Account.is_cash_equivalent
USD = testutil.Amount
@pytest.fixture
def payable_payment_txn():
@ -34,6 +35,14 @@ def payable_payment_txn():
('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):
pred = methodcaller('is_under', *accts)
return data.balance_of(txn, pred)
@ -43,17 +52,18 @@ def test_balance_of_simple_txn():
('Assets:Cash', 50),
('Income:Donations', -50),
])
assert balance_under(txn, 'Assets') == 50
assert balance_under(txn, 'Income') == -50
assert balance_under(txn, 'Assets') == USD(50)
assert balance_under(txn, 'Income') == USD(-50)
def test_zero_balance_of(payable_payment_txn):
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
expected = testutil.Amount(0, '')
assert balance_under(payable_payment_txn, 'Equity') == expected
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):
assert balance_under(payable_payment_txn, 'Assets', 'Expenses') == -50
assert balance_under(payable_payment_txn, 'Assets', 'Liabilities') == -5
assert balance_under(payable_payment_txn, 'Assets', 'Expenses') == USD(-50)
assert balance_under(payable_payment_txn, 'Assets', 'Liabilities') == USD(-5)
def test_multiarg_balance_of():
txn = testutil.Transaction(postings=[
@ -61,7 +71,12 @@ def test_multiarg_balance_of():
('Expenses:BankingFees', 5),
('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):
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'):
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,
currency='USD', cost=None, price=None, flag=None,
**meta):
if not (number is None or isinstance(number, Decimal)):
number = Decimal(number)
if cost is not None:
cost = Cost(*cost)
if meta is None:
meta = None
return bc_data.Posting(
account,
bc_amount.Amount(number, currency),
Amount(number, currency),
cost,
price,
flag,