balance_sheet: Balance.total() accepts multiple account names.

This simplifies the code to total all equity.

conservancy_beancount/reports/balance_sheet.py

@@ -27,6 +27,7 @@ from pathlib import Path
 
 from typing import (
     Any,
+    Collection,
     Dict,
     Hashable,
     Iterable,
@@ -38,6 +39,7 @@ from typing import (
     Sequence,
     TextIO,
     Tuple,
+    Union,
 )
 
 import odf.table  # type:ignore[import]
@@ -127,15 +129,17 @@ class Balances:
             self.balances[key] += post.at_cost()
 
     def total(self,
-              account: Optional[str]=None,
+              account: Union[None, str, Collection[str]]=None,
               classification: Optional[str]=None,
               period: int=Period.ANY,
               fund: int=Fund.ANY,
               post_type: Optional[str]=None,
     ) -> core.Balance:
+        if isinstance(account, str):
+            account = (account,)
         retval = core.MutableBalance()
         for key, balance in self.balances.items():
-            if not (account is None or key.account.is_under(account)):
+            if not (account is None or key.account.is_under(*account)):
                 pass
             elif not (classification is None
                       or key.classification.is_under(classification)):
@@ -328,13 +332,9 @@ class Report(core.BaseODS[Sequence[None], None]):
         self.add_row()
         for fund in [Fund.UNRESTRICTED, Fund.RESTRICTED]:
             preposition = "Without" if fund is Fund.UNRESTRICTED else "With"
-            period_bal = -sum(
-                (self.balances.total(account=account, fund=fund)
-                 for account in EQUITY_ACCOUNTS), core.MutableBalance(),
-            )
-            prior_bal = period_bal + sum(
-                (self.balances.total(account=account, fund=fund, period=Period.PERIOD)
-                 for account in EQUITY_ACCOUNTS), core.MutableBalance(),
+            period_bal = -self.balances.total(account=EQUITY_ACCOUNTS, fund=fund)
+            prior_bal = period_bal + self.balances.total(
+                account=EQUITY_ACCOUNTS, fund=fund, period=Period.PERIOD,
             )
             self.add_row(
                 self.string_cell(f"{preposition} donor restrictions"),
@@ -496,8 +496,7 @@ class Report(core.BaseODS[Sequence[None], None]):
             else:
                 kwargs['period'] = Period.OPENING
         beginnings = [
-            -sum((self.balances.total(account=account, **kwargs)
-                  for account in EQUITY_ACCOUNTS), core.MutableBalance())
+            -self.balances.total(account=EQUITY_ACCOUNTS, **kwargs)
             for kwargs in bal_kwargs
         ]
         self.add_row()
@@ -607,8 +606,7 @@ class Report(core.BaseODS[Sequence[None], None]):
         self.add_row()
 
         totals = [
-            -sum((self.balances.total(account=account, **kwargs)
-                  for account in EQUITY_ACCOUNTS), core.MutableBalance())
+            -self.balances.total(account=EQUITY_ACCOUNTS, **kwargs)
             for kwargs in bal_kwargs
         ]
         self.add_row(

tests/test_reports_balance_sheet.py

@@ -94,6 +94,8 @@ def income_expense_balances():
     ({'fund': Fund.RESTRICTED, 'post_type': 'program'}, 10),
     ({'period': Period.PRIOR, 'fund': Fund.RESTRICTED, 'post_type': 'program'}, '4.80'),
     ({'period': Period.PERIOD, 'fund': Fund.RESTRICTED, 'post_type': 'ø'}, None),
+    ({'account': ('Income', 'Expenses')}, 30),
+    ({'account': ('Income', 'Expenses'), 'fund': Fund.UNRESTRICTED}, 15),
 ])
 def test_balance_total(income_expense_balances, kwargs, expected):
     actual = income_expense_balances.total(**kwargs)