conservancy_beancount/conservancy_beancount/books.py

"""books - Tools for loading the books"""
# 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/>.

import contextlib
import datetime
import os

from pathlib import Path

from beancount import loader as bc_loader

from typing import (
    Any,
    Iterable,
    Iterator,
    Mapping,
    NamedTuple,
    Optional,
    Union,
)
from .beancount_types import (
    LoadResult,
)

PathLike = Union[str, Path]
Year = Union[int, datetime.date]

@contextlib.contextmanager
def workdir(path: PathLike) -> Iterator[Path]:
    old_dir = os.getcwd()
    os.chdir(path)
    try:
        yield Path(old_dir)
    finally:
        os.chdir(old_dir)

class FiscalYear(NamedTuple):
    month: int = 3
    day: int = 1

    def for_date(self, date: Optional[datetime.date]=None) -> int:
        if date is None:
            date = datetime.date.today()
        if (date.month, date.day) < self:
            return date.year - 1
        else:
            return date.year

    def range(self, from_fy: Year, to_fy: Optional[Year]=None) -> Iterable[int]:
        """Return a range of fiscal years

        Both arguments can be either a year (represented as an integer) or a
        date. Dates will be converted into a year by calling for_date() on
        them.

        If the first argument is negative or below 1000, it will be treated as
        an offset. You'll get a range of fiscal years between the second
        argument offset by this amount.

        If the second argument is omitted, it defaults to the current fiscal
        year.

        Note that unlike normal Python ranges, these ranges include the final
        fiscal year.

        Examples:

          range(2015)  # Iterate all fiscal years from 2015 to today, inclusive

          range(-1)  # Iterate the previous fiscal year and current fiscal year
        """
        if not isinstance(from_fy, int):
            from_fy = self.for_date(from_fy)
        if to_fy is None:
            to_fy = self.for_date()
        elif not isinstance(to_fy, int):
            to_fy = self.for_date(to_fy - datetime.timedelta(days=1))
        if from_fy < 1:
            from_fy += to_fy
        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"""

    def __init__(self,
                 books_root: Path,
                 fiscal_year: FiscalYear,
    ) -> 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.
        """
        self.books_root = books_root
        self.opening_root = books_root / 'books'
        self.fiscal_year = fiscal_year

    def _iter_fy_books(self, fy_range: Iterable[int]) -> Iterator[Path]:
        dir_path = self.opening_root
        for year in fy_range:
            path = dir_path / f'{year}.beancount'
            if path.exists():
                yield path
                dir_path = self.books_root

    def fy_range_string(self,
                        from_fy: Year,
                        to_fy: Optional[Year]=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 the books from the first available fiscal
        year through the end of the range.

        Pass the string to Loader.load_string() to actually load data from it.
        """
        paths = self._iter_fy_books(self.fiscal_year.range(from_fy, to_fy))
        try:
            with next(paths).open() as opening_books:
                lines = [opening_books.read()]
        except StopIteration:
            return ''
        for path in paths:
            lines.append(f'include "../{path.name}"')
        return '\n'.join(lines)

    def load_string(self, source: str) -> LoadResult:
        """Load a generated string of Beancount directives

        This method takes a string generated by another Loader method, like
        fy_range_string, and loads it through Beancount, setting up the
        environment as necessary to do that.
        """
        with workdir(self.opening_root):
            retval: LoadResult = bc_loader.load_string(source)
        return retval

    def load_fy_range(self,
                      from_fy: Year,
                      to_fy: Optional[Year]=None,
    ) -> LoadResult:
        """Load books for a range of fiscal years"""
        return self.load_string(self.fy_range_string(from_fy, to_fy))