books: Start Loader class.

conservancy_beancount/books.py

@@ -16,13 +16,18 @@
 
 import datetime
 
+from pathlib import Path
+
 from typing import (
     Iterable,
+    Mapping,
     NamedTuple,
     Optional,
     Union,
 )
 
+PathLike = Union[str, Path]
+PluginsSpec = Mapping[str, Optional[str]]
 Year = Union[int, datetime.date]
 
 class FiscalYear(NamedTuple):
@@ -71,3 +76,74 @@ class FiscalYear(NamedTuple):
         elif from_fy < 1000:
             from_fy, to_fy = to_fy, from_fy + to_fy
         return range(from_fy, to_fy + 1)
+
+
+class Loader:
+    """Load Beancount books organized by fiscal year"""
+
+    DEFAULT_PLUGINS: PluginsSpec = {
+        'conservancy_beancount.plugin': None,
+    }
+
+    def __init__(self,
+                 books_root: Path,
+                 fiscal_year: FiscalYear,
+                 plugins: Optional[PluginsSpec]=None,
+    ) -> None:
+        """Set up a books loader
+
+        Arguments:
+        * books_root: A Path to a Beancount books checkout.
+        * fiscal_year: A FiscalYear object, used to determine what books to
+          load for a given date range.
+        * plugins: A mapping that specifies what plugins should be loaded
+          before any books. The keys are plugin names, and the values are the
+          configuration parameters string to follow. A value of None means the
+          plugin takes no configuration string. By default, the loader loads
+          conservancy_beancount.plugin.
+        """
+        if plugins is None:
+            plugins = self.DEFAULT_PLUGINS
+        self.books_root = books_root
+        self.fiscal_year = fiscal_year
+        self.plugins = dict(plugins)
+
+    def _format_include(self, year: int, subdir: PathLike='') -> str:
+        file_path = Path(self.books_root, subdir, f'{year}.beancount')
+        return f'include "{file_path}"'
+
+    def _format_plugin(self, name: str, optstring: Optional[str]=None) -> str:
+        if optstring is None:
+            return f'plugin "{name}"'
+        else:
+            return f'plugin "{name}" "{optstring}"'
+
+    def fy_range_string(self,
+                        from_fy: Year,
+                        to_fy: Optional[Year]=None,
+                        plugins: Optional[PluginsSpec]=None,
+    ) -> str:
+        """Return a string to load books for a range of fiscal years
+
+        This method generates a range of fiscal years by calling
+        FiscalYear.range() with its first two arguments. It returns a string of
+        Beancount directives to load all plugins and Beancount files for that
+        range of fiscal years, suitable for passing to
+        beancount.loader.load_string().
+
+        You can specify what plugins to load with the plugins argument. If not
+        specified, the string loads the plugins specified for this instance.
+        See the __init__ docstring for details.
+        """
+        if plugins is None:
+            plugins = self.plugins
+        years = iter(self.fiscal_year.range(from_fy, to_fy))
+        try:
+            books_start = self._format_include(next(years), 'books')
+        except StopIteration:
+            return ''
+        return '\n'.join([
+            *(self._format_plugin(name, opts) for name, opts in plugins.items()),
+            books_start,
+            *(self._format_include(year) for year in years),
+        ])

tests/test_books_loader.py

@@ -0,0 +1,76 @@
+"""test_books_loader - Unit tests for books Loader class"""
+# Copyright © 2020  Brett Smith
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+
+from datetime import date
+from pathlib import Path
+
+import pytest
+
+from . import testutil
+
+from conservancy_beancount import books
+
+books_path = testutil.test_path('booksroot')
+
+@pytest.fixture(scope='module')
+def conservancy_loader():
+    return books.Loader(books_path, books.FiscalYear(3))
+
+def format_include(year, subdir=''):
+    path = Path(books_path, subdir, f'{year}.beancount')
+    return f'include "{path}"'
+
+def format_plugin(name, optstring=None):
+    if optstring is None:
+        return f'plugin "{name}"'
+    else:
+        return f'plugin "{name}" "{optstring}"'
+
+def expect_string(years, plugins={'conservancy_beancount.plugin': None}):
+    years = iter(years)
+    year1_s = format_include(next(years), 'books')
+    return '\n'.join([
+        *(format_plugin(name, opts) for name, opts in plugins.items()),
+        year1_s,
+        *(format_include(year) for year in years),
+    ])
+
+@pytest.mark.parametrize('range_start,range_stop,expect_years', [
+    (2019, 2020, [2019, 2020]),
+    (-1, 2020, [2019, 2020]),
+    (date(2019, 1, 1), date(2020, 6, 1), range(2018, 2021)),
+    (-1, date(2020, 2, 1), [2018, 2019]),
+])
+def test_fy_range_string(conservancy_loader, range_start, range_stop, expect_years):
+    expected = expect_string(expect_years)
+    assert conservancy_loader.fy_range_string(range_start, range_stop) == expected
+
+@pytest.mark.parametrize('year_offset', range(-3, 1))
+def test_fy_range_string_one_offset(conservancy_loader, year_offset):
+    this_year = date.today().year
+    expected = expect_string(range(this_year + year_offset, this_year + 1))
+    assert conservancy_loader.fy_range_string(year_offset) == expected
+
+@pytest.mark.parametrize('plugins', [
+    {},
+    {'conservancy_beancount.plugin': '-all'},
+])
+def test_fy_range_string_plugins_override(conservancy_loader, plugins):
+    expected = expect_string([2019, 2020], plugins)
+    assert conservancy_loader.fy_range_string(2019, 2020, plugins) == expected
+
+def test_fy_range_string_empty_range(conservancy_loader):
+    assert conservancy_loader.fy_range_string(2020, 2019) == ''