From fd3bd683267382726e41cef4ffd4ce2bc39d2a5f Mon Sep 17 00:00:00 2001
From: Brett Smith <brettcsmith@brettcsmith.org>
Date: Thu, 16 Jul 2020 13:51:23 -0400
Subject: [PATCH] data: Add Account iteration methods.

---
 conservancy_beancount/data.py | 33 +++++++++++++++++++
 tests/test_data_account.py    | 60 +++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/conservancy_beancount/data.py b/conservancy_beancount/data.py
index 55c995d..33b036d 100644
--- a/conservancy_beancount/data.py
+++ b/conservancy_beancount/data.py
@@ -201,6 +201,39 @@ class Account(str):
     def is_account(cls, s: str) -> bool:
         return cls.ACCOUNT_RE.fullmatch(s) is not None
 
+    @classmethod
+    def iter_accounts_by_classification(cls, s: str) -> Iterator['Account']:
+        class_re = re.compile(f'^{re.escape(s)}(?::|$)')
+        for name, meta in cls._meta_map.items():
+            try:
+                match = class_re.match(meta['classification'])
+            except KeyError:
+                match = None
+            if match:
+                yield cls(name)
+
+    @classmethod
+    def iter_accounts_by_hierarchy(cls, s: str) -> Iterator['Account']:
+        for name in cls._meta_map:
+            account = cls(name)
+            if account.is_under(s):
+                yield account
+
+    @classmethod
+    def iter_accounts(cls, s: str) -> Iterator['Account']:
+        """Iterate account objects by name or classification
+
+        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.
+        """
+        # We append a stub subaccount to match root accounts.
+        if cls.is_account(f'{s}:Test'):
+            return cls.iter_accounts_by_hierarchy(s)
+        else:
+            return cls.iter_accounts_by_classification(s)
+
     @property
     def meta(self) -> AccountMeta:
         return self._meta_map[self]
diff --git a/tests/test_data_account.py b/tests/test_data_account.py
index aed5909..5c78d1d 100644
--- a/tests/test_data_account.py
+++ b/tests/test_data_account.py
@@ -27,6 +27,25 @@ from conservancy_beancount import data
 
 clean_account_meta = pytest.fixture()(testutil.clean_account_meta)
 
+@pytest.fixture
+def asset_hierarchy():
+    entries = [
+        Open({'classification': 'Investment'},
+             Date(2002, 2, 1), 'Assets:Bank:CD', None, None),
+        Open({'classification': 'Cash'},
+             Date(2002, 2, 1), 'Assets:Bank:Checking', None, None),
+        Open({'classification': 'Cash'},
+             Date(2002, 2, 1), 'Assets:Bank:Savings', None, None),
+        Open({'classification': 'Cash'},
+             Date(2002, 2, 1), 'Assets:Cash', None, None),
+        Open({'classification': 'Investment'},
+             Date(2002, 2, 1), 'Assets:Investment:Commodities', None, None),
+        Open({'classification': 'Investment'},
+             Date(2002, 2, 1), 'Assets:Investment:Stocks', None, None),
+    ]
+    data.Account.load_openings_and_closings(entries)
+    yield from testutil.clean_account_meta()
+
 def check_account_meta(acct_meta, opening, closing=None):
     if isinstance(acct_meta, str):
         acct_meta = data.Account(acct_meta).meta
@@ -333,3 +352,44 @@ def test_load_from_books(clean_account_meta):
     check_meta = data.Account(entries[0].account).meta
     assert check_meta.open_date == entries[0].date
     assert check_meta.close_date == entries[-1].date
+
+@pytest.mark.parametrize('arg,expect_subaccts', [
+    ('Assets', ['Bank:CD', 'Bank:Checking', 'Bank:Savings', 'Cash',
+                'Investment:Commodities', 'Investment:Stocks']),
+    ('Assets:Bank', ['CD', 'Checking', 'Savings']),
+    ('Assets:Investment', ['Commodities', 'Stocks']),
+    ('Equity', []),
+])
+def test_iter_accounts_by_hierarchy(asset_hierarchy, arg, expect_subaccts):
+    assert set(data.Account.iter_accounts_by_hierarchy(arg)) == {
+        f'{arg}:{sub}' for sub in expect_subaccts
+    }
+
+@pytest.mark.parametrize('arg,expect_subaccts', [
+    ('Cash', ['Bank:Checking', 'Bank:Savings', 'Cash']),
+    ('Investment', ['Bank:CD', 'Investment:Commodities', 'Investment:Stocks']),
+    ('Equity', []),
+])
+def test_iter_accounts_by_classification(asset_hierarchy, arg, expect_subaccts):
+    assert set(data.Account.iter_accounts_by_classification(arg)) == {
+        f'Assets:{sub}' for sub in expect_subaccts
+    }
+
+@pytest.mark.parametrize('arg,expect_subaccts', [
+    ('Assets', ['Bank:CD', 'Bank:Checking', 'Bank:Savings', 'Cash',
+                'Investment:Commodities', 'Investment:Stocks']),
+    ('Assets:Bank', ['CD', 'Checking', 'Savings']),
+    ('Assets:Investment', ['Commodities', 'Stocks']),
+    ('Cash', ['Bank:Checking', 'Bank:Savings', 'Cash']),
+    ('Investment', ['Bank:CD', 'Investment:Commodities', 'Investment:Stocks']),
+    ('Equity', []),
+    ('Unused classification', []),
+])
+def test_iter_accounts(asset_hierarchy, arg, expect_subaccts):
+    if arg.startswith('Assets'):
+        prefix = arg
+    else:
+        prefix = 'Assets'
+    assert set(data.Account.iter_accounts(arg)) == {
+        f'{prefix}:{sub}' for sub in expect_subaccts
+    }