fund: Use Balances instead of PeriodPostings.

A few motivations for this:

* This makes the fund report more maintainable, because the data structure
  is better suited to the task at hand, making it easier to follow what's
  going on.

* This helps put Balances through a little more testing with a high-level
  report.

* This makes the fund report a little faster, since totals are usually
  calculated from a smaller number of Balance objects rather than all
  Postings over a period.

conservancy_beancount/reports/core.py

@@ -467,6 +467,12 @@ class Balances:
                  or account.meta.close_date > start_date)
         )
 
+    def meta_values(self) -> Set[str]:
+        retval = {key.post_meta for key in self.balances}
+        retval.discard(None)
+        # discarding None ensures we return the desired type.
+        return retval  # type:ignore[return-value]
+
 
 class RelatedPostings(Sequence[data.Posting]):
     """Collect and query related postings

conservancy_beancount/reports/fund.py

@@ -47,6 +47,7 @@ import argparse
 import collections
 import datetime
 import enum
+import functools
 import locale
 import logging
 import sys
@@ -80,8 +81,7 @@ from .. import cliutil
 from .. import config as configmod
 from .. import data
 
-AccountsMap = Mapping[data.Account, core.PeriodPostings]
-FundPosts = Tuple[MetaValue, AccountsMap]
+Period = core.Period
 
 EQUITY_ACCOUNTS = ['Income', 'Expenses', 'Equity']
 INFO_ACCOUNTS = [
@@ -94,25 +94,34 @@ PROGNAME = 'fund-report'
 UNRESTRICTED_FUND = 'Conservancy'
 logger = logging.getLogger('conservancy_beancount.reports.fund')
 
-class ODSReport(core.BaseODS[FundPosts, None]):
-    def __init__(self, start_date: datetime.date, stop_date: datetime.date) -> None:
+class ODSReport(core.BaseODS[str, None]):
+    def __init__(self,
+                 balances: core.Balances,
+                 start_date: datetime.date,
+                 stop_date: datetime.date,
+    ) -> None:
         super().__init__()
+        self.balances = balances
         self.start_date = start_date
         self.stop_date = stop_date
 
-    def section_key(self, row: FundPosts) -> None:
+    def section_key(self, row: str) -> None:
         return None
 
-    def start_spreadsheet(self, *, expanded: bool=True) -> None:
+    def start_spreadsheet(self) -> None:
         headers = [["Fund"], ["Balance as of", self.start_date.isoformat()]]
-        if expanded:
+        # If self.sheet has children, that means we've already written the
+        # first sheet and we're starting the second, expanded sheet.
+        if self.sheet.childNodes:
+            self.expanded = True
             sheet_name = "With Breakdowns"
             headers.extend([acct] for acct in EQUITY_ACCOUNTS)
         else:
+            self.expanded = False
             sheet_name = "Fund Report"
             headers += [["Additions"], ["Releases from", "Restrictions"]]
         headers.append(["Balance as of", self.stop_date.isoformat()])
-        if expanded:
+        if self.expanded:
             headers += [
                 ["Of which", "Receivable"],
                 ["Of which", "Prepaid Expenses"],
@@ -151,108 +160,130 @@ class ODSReport(core.BaseODS[FundPosts, None]):
             numbercolumnsspanned=6,
         ))
         self.add_row()
+        self.sheet_totals = [core.MutableBalance() for _ in range(len(headers) - 1)]
 
-    def end_spreadsheet(self) -> None:
-        start_sheet = self.sheet
-        self.set_open_sheet(self.sheet)
-        self.start_spreadsheet(expanded=False)
-        bal_indexes = [0, 1, 2, 4]
-        totals = [core.MutableBalance() for _ in bal_indexes]
-        threshold = Decimal('.5')
-        for fund, source_bals in self.balances.items():
-            balances = [source_bals[index] for index in bal_indexes]
-            # Incorporate Equity changes to Release from Restrictions.
-            # Note that using -= mutates the balance in a way we don't want.
-            balances[2] = balances[2] - source_bals[3]
-            if (not all(bal.clean_copy(threshold).le_zero() for bal in balances)
-                and fund != UNRESTRICTED_FUND):
-                self.write_balances(fund, balances)
-                for total, bal in zip(totals, balances):
-                    total += bal
-        self.write_balances('', totals, self.style_bottomline)
-        self.document.spreadsheet.childNodes.reverse()
-        self.sheet = start_sheet
-
-    def _row_balances(self, accounts_map: AccountsMap) -> Iterator[core.Balance]:
-        key_order = [core.OPENING_BALANCE_NAME, *EQUITY_ACCOUNTS, core.ENDING_BALANCE_NAME]
-        balances: Dict[str, core.Balance] = {key: core.MutableBalance() for key in key_order}
-        for acct_s, balance in core.account_balances(accounts_map, EQUITY_ACCOUNTS):
-            if acct_s in balances:
-                balances[acct_s] = balance
-            else:
-                acct_root, _, _ = acct_s.partition(':')
-                balances[acct_root] += balance
-        for key in key_order:
-            if key == 'Expenses':
-                yield balances[key]
-            else:
-                yield -balances[key]
-        for info_key in INFO_ACCOUNTS:
-            for _, balance in core.account_balances(accounts_map, [info_key]):
-                pass
-            yield core.normalize_amount_func(info_key)(balance)
-
-    def write_balances(self,
-                       fund: str,
-                       balances: Iterable[core.Balance],
-                       style: Union[None, str, odf.style.Style]=None,
+    def write_balance_row(self,
+                          name: str,
+                          balances: Iterable[core.Balance],
+                          style: Union[None, str, odf.style.Style]=None,
     ) -> odf.table.TableRow:
         return self.add_row(
-            self.string_cell(fund, stylename=self.style_endtext),
+            self.string_cell(name, stylename=self.style_endtext),
             *(self.balance_cell(bal, stylename=style) for bal in balances),
         )
 
-    def write_row(self, row: FundPosts) -> None:
-        fund, accounts_map = row
-        self.balances[fund] = list(self._row_balances(accounts_map))
-        if fund != UNRESTRICTED_FUND:
-            self.write_balances(fund, self.balances[fund])
-
-    def write(self, rows: Iterable[FundPosts]) -> None:
-        self.balances: Dict[str, Sequence[core.Balance]] = collections.OrderedDict()
-        super().write(rows)
-        try:
-            unrestricted = self.balances[UNRESTRICTED_FUND]
-        except KeyError:
-            pass
+    def write_balances(self,
+                       name: str,
+                       style: Union[None, str, odf.style.Style]=None,
+                       fund: int=core.Fund.ANY,
+    ) -> odf.table.TableRow:
+        if (fund & core.Fund.ANY) == core.Fund.ANY:
+            total = functools.partial(self.balances.total, post_meta=name)
         else:
-            self.add_row()
-            self.write_balances("Unrestricted", unrestricted)
+            total = functools.partial(self.balances.total, fund=fund)
+        balances = [
+            -total(account=EQUITY_ACCOUNTS, period=Period.THRU_MIDDLE),
+            -total(account='Income', period=Period.PERIOD),
+            total(account='Expenses', period=Period.PERIOD),
+            -total(account='Equity', period=Period.PERIOD),
+        ]
+        if not self.expanded:
+            threshold = Decimal('.5')
+            if all(bal.clean_copy(threshold).le_zero() for bal in balances):
+                return
+            equity = balances.pop()
+            balances[-1] -= equity
+        balances.append(-total(account=EQUITY_ACCOUNTS))
+        if self.expanded:
+            for account in INFO_ACCOUNTS:
+                norm_func = core.normalize_amount_func(account)
+                balances.append(norm_func(total(account=account)))
+        for sheet_tot, row_tot in zip(self.sheet_totals, balances):
+            sheet_tot += row_tot
+        return self.write_balance_row(name, balances, style)
+
+    def write_row(self, row: str) -> None:
+        if row != UNRESTRICTED_FUND:
+            self.write_balances(row)
+
+    def write(self, rows: Iterable[str]) -> None:
+        row_list = list(rows)
+        # Write the basic, auditor-style fund report.
+        super().write(iter(row_list))
+        self.write_balance_row("", self.sheet_totals, self.style_bottomline)
+        # Write the expanded fund report. start_spreadsheet() will see we've
+        # written the first sheet and adapt.
+        super().write(iter(row_list))
+        self.write_balances("Unrestricted", fund=core.Fund.UNRESTRICTED)
+        self.set_open_sheet(self.sheet)
 
 
 class TextReport:
     def __init__(self,
+                 balances: core.Balances,
                  start_date: datetime.date,
                  stop_date: datetime.date,
                  out_file: TextIO) -> None:
+        self.balances = balances
         self.start_date = start_date
         self.stop_date = stop_date
         self.out_file = out_file
+        accounts = self.balances.iter_accounts()
+        self.equity_accounts = [
+            account for _, account in
+            core.sort_and_filter_accounts(accounts, EQUITY_ACCOUNTS)
+        ]
+        self.info_accounts = [
+            account for _, account in
+            core.sort_and_filter_accounts(accounts, INFO_ACCOUNTS)
+        ]
 
-    def _account_balances(self,
-                          fund: str,
-                          account_map: AccountsMap,
-    ) -> Iterator[Tuple[str, Sequence[str]]]:
+    def _format_total(self,
+                      fund_name: str,
+                      period: int=Period.ANY,
+                      account: Optional[str]=None,
+    ) -> Sequence[str]:
+        if account is None:
+            account_arg: Union[str, Sequence[str]] = EQUITY_ACCOUNTS
+            account_exact = False
+        else:
+            account_arg = account
+            account_exact = True
+        balance = self.balances.total(
+            post_meta=fund_name,
+            period=period,
+            account=account_arg,
+            account_exact=account_exact,
+        )
+        if balance.is_zero():
+            return []
+        elif account is None:
+            balance = -balance
+        else:
+            norm_func = core.normalize_amount_func(f'{account}:RootsOK')
+            balance = norm_func(balance)
+        return balance.format(None, sep='\0').split('\0')
+
+    def _account_balances(self, fund: str) -> Iterator[Tuple[str, Sequence[str]]]:
         total_fmt = f'{fund} balance as of {{}}'
-        for acct_s, balance in core.account_balances(account_map, EQUITY_ACCOUNTS):
-            if acct_s is core.OPENING_BALANCE_NAME:
-                acct_s = total_fmt.format(self.start_date.isoformat())
-            elif acct_s is core.ENDING_BALANCE_NAME:
-                acct_s = total_fmt.format(self.stop_date.isoformat())
-            if not acct_s.startswith('Expenses:'):
-                balance = -balance
-            yield acct_s, balance.format(None, sep='\0').split('\0')
-        for _, account in core.sort_and_filter_accounts(account_map, INFO_ACCOUNTS):
-            balance = account_map[account].stop_bal
-            if not balance.is_zero():
-                balance = core.normalize_amount_func(account)(balance)
-                yield account, balance.format(None, sep='\0').split('\0')
+        yield (total_fmt.format(self.start_date.isoformat()),
+               self._format_total(fund, Period.THRU_MIDDLE))
+        for account in self.equity_accounts:
+            total = self._format_total(fund, Period.PERIOD, account)
+            if total:
+                yield (account, total)
+        yield (total_fmt.format(self.stop_date.isoformat()),
+               self._format_total(fund))
+        for account in self.info_accounts:
+            total = self._format_total(fund, account=account)
+            if total:
+                yield (account, total)
 
-    def write(self, rows: Iterable[FundPosts]) -> None:
+    def write(self, rows: Iterable[str]) -> None:
         output = [
             line
-            for fund, account_map in rows
-            for line in self._account_balances(fund, account_map)
+            for fund in rows
+            for line in self._account_balances(fund)
         ]
         acct_width = max(len(acct_s) for acct_s, _ in output) + 2
         bal_width = max(len(s) for _, bal_s in output for s in bal_s)
@@ -369,10 +400,10 @@ def main(arglist: Optional[Sequence[str]]=None,
     returncode = 0
     books_loader = config.books_loader()
     if books_loader is None:
-        entries, load_errors, _ = books.Loader.load_none(config.config_file_path())
+        entries, load_errors, options = books.Loader.load_none(config.config_file_path())
         returncode = cliutil.ExitCode.NoConfiguration
     else:
-        entries, load_errors, _ = books_loader.load_fy_range(args.start_date, args.stop_date)
+        entries, load_errors, options = books_loader.load_fy_range(args.start_date, args.stop_date)
         if load_errors:
             returncode = cliutil.ExitCode.BeancountErrors
         elif not entries:
@@ -380,6 +411,7 @@ def main(arglist: Optional[Sequence[str]]=None,
     for error in load_errors:
         bc_printer.print_error(error, file=stderr)
 
+    data.Account.load_from_books(entries, options)
     postings = iter(
         post
         for post in data.Posting.from_entries(entries)
@@ -395,27 +427,19 @@ def main(arglist: Optional[Sequence[str]]=None,
         postings = ruleset.rewrite(postings)
     for search_term in args.search_terms:
         postings = search_term.filter_postings(postings)
-    fund_postings = {
-        key: related
-        for key, related in core.RelatedPostings.group_by_meta(postings, 'project')
-        if isinstance(key, str)
-    }
-    period_cls = core.PeriodPostings.with_start_date(args.start_date)
-    fund_map = collections.OrderedDict(
-        (fund, dict(period_cls.group_by_account(fund_postings[fund])))
-        for fund in sorted(fund_postings, key=lambda s: locale.strxfrm(s.casefold()))
-    )
-    if not fund_map:
+    balances = core.Balances(postings, args.start_date, args.stop_date, 'project')
+    funds = sorted(balances.meta_values(), key=lambda s: locale.strxfrm(s.casefold()))
+    if not funds:
         logger.warning("no matching postings found to report")
         returncode = returncode or cliutil.ExitCode.NoDataFiltered
     elif args.report_type is ReportType.TEXT:
         out_file = cliutil.text_output(args.output_file, stdout)
-        report = TextReport(args.start_date, args.stop_date, out_file)
-        report.write(fund_map.items())
+        report = TextReport(balances, args.start_date, args.stop_date, out_file)
+        report.write(funds)
     else:
-        ods_report = ODSReport(args.start_date, args.stop_date)
+        ods_report = ODSReport(balances, args.start_date, args.stop_date)
         ods_report.set_common_properties(config.books_repo())
-        ods_report.write(fund_map.items())
+        ods_report.write(funds)
         if args.output_file is None:
             out_dir_path = config.repository_path() or Path()
             args.output_file = out_dir_path / 'FundReport_{}_{}.ods'.format(

tests/test_reports_fund.py

@@ -93,6 +93,8 @@ BALANCES_BY_YEAR = {
     ],
 }
 
+clean_account_meta = pytest.fixture(autouse=True)(testutil.clean_account_meta)
+
 @pytest.fixture
 def fund_entries():
     return copy.deepcopy(_ledger_load[0])