From aaa26e9e619c9ca6dec8797154757e8d2448aff6 Mon Sep 17 00:00:00 2001 From: Brett Smith Date: Mon, 20 Jul 2020 11:22:23 -0400 Subject: [PATCH] data: Make Account.iter_accounts() argument optional. --- conservancy_beancount/data.py | 7 +++++-- tests/test_data_account.py | 4 +++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/conservancy_beancount/data.py b/conservancy_beancount/data.py index 829cec7..b102f3b 100644 --- a/conservancy_beancount/data.py +++ b/conservancy_beancount/data.py @@ -220,16 +220,19 @@ class Account(str): yield account @classmethod - def iter_accounts(cls, s: str) -> Iterator['Account']: + def iter_accounts(cls, s: Optional[str]=None) -> Iterator['Account']: """Iterate account objects by name or classification + With no argument, returns an iterator of all known account names. If you pass in a root account name, or a valid account string, returns an iterator of all accounts under that account in the hierarchy. Otherwise, returns an iterator of all accounts with the given ``classification`` metadata. """ + if s is None: + return (cls(acct) for acct in cls._meta_map) # We append a stub subaccount to match root accounts. - if cls.is_account(f'{s}:Test'): + elif cls.is_account(f'{s}:RootsOK'): return cls.iter_accounts_by_hierarchy(s) else: return cls.iter_accounts_by_classification(s) diff --git a/tests/test_data_account.py b/tests/test_data_account.py index 9546d32..106d114 100644 --- a/tests/test_data_account.py +++ b/tests/test_data_account.py @@ -444,6 +444,8 @@ def test_iter_accounts_by_classification(asset_hierarchy, arg, expect_subaccts): } @pytest.mark.parametrize('arg,expect_subaccts', [ + (None, ['Bank:CD', 'Bank:Checking', 'Bank:Savings', 'Cash', + 'Investment:Commodities', 'Investment:Stocks']), ('Assets', ['Bank:CD', 'Bank:Checking', 'Bank:Savings', 'Cash', 'Investment:Commodities', 'Investment:Stocks']), ('Assets:Bank', ['CD', 'Checking', 'Savings']), @@ -454,7 +456,7 @@ def test_iter_accounts_by_classification(asset_hierarchy, arg, expect_subaccts): ('Unused classification', []), ]) def test_iter_accounts(asset_hierarchy, arg, expect_subaccts): - if arg.startswith('Assets'): + if arg and arg.startswith('Assets'): prefix = arg else: prefix = 'Assets'