From 9595d3334d5bc609335b2a1d58da565af9675c1f Mon Sep 17 00:00:00 2001
From: Brett Smith <brettcsmith@brettcsmith.org>
Date: Mon, 25 May 2020 10:37:21 -0400
Subject: [PATCH] books: Add Loader.load_all method.

---
 conservancy_beancount/books.py | 45 ++++++++++++++++++++++------------
 tests/test_books_loader.py     | 38 +++++++++++++++++++---------
 2 files changed, 57 insertions(+), 26 deletions(-)

diff --git a/conservancy_beancount/books.py b/conservancy_beancount/books.py
index 002f7ad..e3fac9a 100644
--- a/conservancy_beancount/books.py
+++ b/conservancy_beancount/books.py
@@ -118,6 +118,35 @@ class Loader:
             if path.exists():
                 yield path
 
+    def _load_paths(self, paths: Iterator[Path]) -> LoadResult:
+        try:
+            entries, errors, options_map = bc_loader.load_file(next(paths))
+        except StopIteration:
+            entries, errors, options_map = [], [], {}
+        for load_path in paths:
+            new_entries, new_errors, new_options = bc_loader.load_file(load_path)
+            # We only want transactions from the new fiscal year.
+            # We don't want the opening balance, duplicate definitions, etc.
+            fy_filename = str(load_path.parent.parent / load_path.name)
+            entries.extend(
+                entry for entry in new_entries
+                if entry.meta.get('filename') == fy_filename
+            )
+            errors.extend(new_errors)
+        return entries, errors, options_map
+
+    def load_all(self) -> LoadResult:
+        """Load all of the books
+
+        This method loads all of the books. It finds the books by simply
+        globbing the filesystem. It still loads each fiscal year in sequence to
+        provide the best cache utilization.
+        """
+        path = Path(self.books_root, 'books')
+        fy_paths = list(path.glob('[1-9][0-9][0-9][0-9].beancount'))
+        fy_paths.sort()
+        return self._load_paths(iter(fy_paths))
+
     def load_fy_range(self,
                       from_fy: Year,
                       to_fy: Optional[Year]=None,
@@ -131,18 +160,4 @@ class Loader:
         """
         fy_range = self.fiscal_year.range(from_fy, to_fy)
         fy_paths = self._iter_fy_books(fy_range)
-        try:
-            entries, errors, options_map = bc_loader.load_file(next(fy_paths))
-        except StopIteration:
-            entries, errors, options_map = [], [], {}
-        for load_path in fy_paths:
-            new_entries, new_errors, new_options = bc_loader.load_file(load_path)
-            # We only want transactions from the new fiscal year.
-            # We don't want the opening balance, duplicate definitions, etc.
-            fy_filename = str(load_path.parent.parent / load_path.name)
-            entries.extend(
-                entry for entry in new_entries
-                if entry.meta.get('filename') == fy_filename
-            )
-            errors.extend(new_errors)
-        return entries, errors, options_map
+        return self._load_paths(fy_paths)
diff --git a/tests/test_books_loader.py b/tests/test_books_loader.py
index 86c764f..ce6ae87 100644
--- a/tests/test_books_loader.py
+++ b/tests/test_books_loader.py
@@ -14,6 +14,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+import collections
 import re
 
 from datetime import date
@@ -23,6 +24,7 @@ import pytest
 
 from . import testutil
 
+from beancount.core import data as bc_data
 from conservancy_beancount import books
 
 books_path = testutil.test_path('books')
@@ -31,6 +33,20 @@ books_path = testutil.test_path('books')
 def conservancy_loader():
     return books.Loader(books_path, books.FiscalYear(3))
 
+def check_openings(entries):
+    openings = collections.defaultdict(int)
+    for entry in entries:
+        if isinstance(entry, bc_data.Open):
+            openings[entry.account] += 1
+    for account, count in openings.items():
+        assert count == 1, f"found {count} open directives for {account}"
+
+def get_narrations(entries):
+    return {
+        entry.narration for entry in entries
+        if isinstance(entry, bc_data.Transaction)
+    }
+
 @pytest.mark.parametrize('from_fy,to_fy,expect_years', [
     (2019, 2019, range(2019, 2020)),
     (0, 2019, range(2019, 2020)),
@@ -47,26 +63,26 @@ def conservancy_loader():
 def test_load_fy_range(conservancy_loader, from_fy, to_fy, expect_years):
     entries, errors, options_map = conservancy_loader.load_fy_range(from_fy, to_fy)
     assert not errors
-    narrations = {getattr(entry, 'narration', None) for entry in entries}
+    narrations = get_narrations(entries)
     assert ('2018 donation' in narrations) == (2018 in expect_years)
     assert ('2019 donation' in narrations) == (2019 in expect_years)
     assert ('2020 donation' in narrations) == (2020 in expect_years)
 
 def test_load_fy_range_does_not_duplicate_openings(conservancy_loader):
     entries, errors, options_map = conservancy_loader.load_fy_range(2010, 2030)
-    openings = []
-    open_accounts = set()
-    for entry in entries:
-        try:
-            open_accounts.add(entry.account)
-        except AttributeError:
-            pass
-        else:
-            openings.append(entry)
-    assert len(openings) == len(open_accounts)
+    check_openings(entries)
 
 def test_load_fy_range_empty(conservancy_loader):
     entries, errors, options_map = conservancy_loader.load_fy_range(2020, 2019)
     assert not errors
     assert not entries
     assert not options_map
+
+def test_load_all(conservancy_loader):
+    entries, errors, options_map = conservancy_loader.load_all()
+    assert not errors
+    narrations = get_narrations(entries)
+    assert '2018 donation' in narrations
+    assert '2019 donation' in narrations
+    assert '2020 donation' in narrations
+    check_openings(entries)