diff --git a/conservancy_beancount/reports/balance_sheet.py b/conservancy_beancount/reports/balance_sheet.py index 14a5caf..f937b72 100644 --- a/conservancy_beancount/reports/balance_sheet.py +++ b/conservancy_beancount/reports/balance_sheet.py @@ -148,12 +148,25 @@ class Balances: period: int=Period.ANY, fund: int=Fund.ANY, post_type: Optional[str]=None, + *, + account_exact: bool=False, ) -> core.Balance: if isinstance(account, str): account = (account,) + acct_pred: Callable[[data.Account], bool] + if account is None: + acct_pred = lambda acct: True + elif account_exact: + # At this point, between this isinstance() above and the earlier + # `account is None` check, we've collapsed the type of `account` to + # `Collection[str]`. Unfortunately the logic is too involved for + # mypy to follow, so ignore the type problem. + acct_pred = lambda acct: acct in account # type:ignore[operator] + else: + acct_pred = lambda acct: acct.is_under(*account) is not None # type:ignore[misc] retval = core.MutableBalance() for key, balance in self.balances.items(): - if not (account is None or key.account.is_under(*account)): + if not acct_pred(key.account): pass elif not (classification is None or key.classification.is_under(classification)): @@ -599,10 +612,14 @@ class Report(core.BaseODS[Sequence[None], None]): want_balance = acct_root not in EQUITY_ACCOUNTS self.add_row() for account in self.balances.iter_accounts(acct_root): - period_bal = self.balances.total(account=account, period=Period.PERIOD) - prior_bal = self.balances.total(account=account, period=Period.PRIOR) + period_bal = self.balances.total( + account=account, period=Period.PERIOD, account_exact=True, + ) + prior_bal = self.balances.total( + account=account, period=Period.PRIOR, account_exact=True, + ) if want_balance: - close_bal = self.balances.total(account=account) + close_bal = self.balances.total(account=account, account_exact=True) close_cell = self.balance_cell(norm_func(close_bal)) open_cell = self.balance_cell(norm_func(close_bal - period_bal)) else: