conservancy_beancount/conservancy_beancount/tools/opening_balances.py

#!/usr/bin/env python3
"""opening_balances.py - Tool to generate opening balances transactions

This tool generates an opening balances transaction for a given date and writes
it to stdout. Use this when you close the books for a year to record the final
balances for that year.

Run it without arguments to generate opening balances for the current fiscal
year. You can also specify a fiscal year to generate opening balances for, or
even a specific date (which can be helpful for testing or debugging).
"""
# SPDX-FileCopyrightText: © 2020 Martin Michlmayr <tbm@cyrius.com>
# SPDX-FileCopyrightText: © 2020 Brett Smith
# License: AGPLv3-or-later WITH Beancount-Plugin-Additional-Permission-1.0
#
# Full copyright and licensing details can be found at toplevel file
# LICENSE.txt in the repository.

import argparse
import collections
import copy
import datetime
import enum
import locale
import logging
import sys

from typing import (
    Dict,
    Iterable,
    Iterator,
    Mapping,
    NamedTuple,
    Optional,
    Sequence,
    TextIO,
    Tuple,
)
from ..beancount_types import (
    Error,
    MetaKey,
    MetaValue,
    Sortable,
    Transaction,
)

from decimal import Decimal, ROUND_HALF_EVEN, ROUND_HALF_UP

from .. import books
from .. import cliutil
from .. import config as configmod
from .. import data
from ..reports.core import Balance

from beancount.core import data as bc_data
from beancount.core import display_context as bc_dcontext
from beancount.parser import printer as bc_printer

from beancount.core.convert import get_cost
from beancount.core.inventory import Inventory
from beancount.core.position import Position, get_position

RESTRICTED_ACCOUNT = data.Account('Equity:Funds:Restricted')
UNRESTRICTED_ACCOUNT = data.Account('Equity:Funds:Unrestricted')
PROGNAME = 'opening-balances'
logger = logging.getLogger('conservancy_beancount.tools.opening_balances')

def quantize_amount(
        amount: data.Amount,
        exp: Decimal=Decimal('.01'),
        rounding: str=ROUND_HALF_EVEN,
) -> data.Amount:
    return amount._replace(number=amount.number.quantize(exp, rounding=rounding))

class AccountWithFund(NamedTuple):
    account: data.Account
    fund: Optional[MetaValue]

    def sortkey(self) -> Sortable:
        account, fund = self
        return (
            0 if fund is None else 1,
            locale.strxfrm(account),
            locale.strxfrm(str(fund).casefold()),
        )


class Posting(data.Posting):
    @staticmethod
    def _position_sortkey(position: Position) -> str:
        units, cost = position
        if cost is None:
            # Beancount type-declares that position.cost must be a Cost, but
            # in practice that's not true. Call get_position(post) on any
            # post without a cost and see what it returns. Hence the ignore.
            return units.currency  # type:ignore[unreachable]
        else:
            return f'{units.currency} {cost.currency} {cost.date.isoformat()}'

    @classmethod
    def build_opening(
            cls,
            key: AccountWithFund,
            meta_key: MetaKey,
            inventory: Inventory,
    ) -> Iterator[bc_data.Posting]:
        account, project = key
        if project is None:
            meta: Optional[Dict[MetaKey, MetaValue]] = None
        else:
            meta = {meta_key: project}
        for units, cost in sorted(inventory, key=cls._position_sortkey):
            if cost is None:
                units = quantize_amount(units)
            yield bc_data.Posting(
                account, units, cost, None, None, copy.copy(meta),
            )


def parse_arguments(arglist: Optional[Sequence[str]]=None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(prog=PROGNAME)
    cliutil.add_version_argument(parser)
    cliutil.add_loglevel_argument(parser)
    parser.add_argument(
        '--fund-metadata-key', '-m',
        metavar='KEY',
        dest='meta_key',
        default='project',
        help="""Name of the fund metadata key. Default %(default)s.
""")
    parser.add_argument(
        '--unrestricted-fund', '-u',
        metavar='PROJECT',
        default='Conservancy',
        help="""Name of the unrestricted fund. Default %(default)s.
""")
    parser.add_argument(
        'as_of_date',
        metavar='YEAR_OR_DATE',
        type=cliutil.year_or_date_arg,
        nargs='?',
        help="""Date to generate opening balances for. You can provide just
a year to generate balances for the start of that fiscal year. Defaults to the
current fiscal year.
""")
    return parser.parse_args(arglist)

def main(arglist: Optional[Sequence[str]]=None,
         stdout: TextIO=sys.stdout,
         stderr: TextIO=sys.stderr,
         config: Optional[configmod.Config]=None,
) -> int:
    args = parse_arguments(arglist)
    cliutil.set_loglevel(logger, args.loglevel)
    if config is None:
        config = configmod.Config()
        config.load_file()

    fy = config.fiscal_year_begin()
    if args.as_of_date is None:
        args.as_of_date = fy.for_date()
    if isinstance(args.as_of_date, int):
        args.as_of_date = fy.first_date(args.as_of_date)

    books_load = books.Loader.dispatch(config.books_loader(), 0, args.as_of_date)
    returncode = books_load.returncode()
    books_load.print_errors(stderr)

    inventories: Mapping[AccountWithFund, Inventory] = collections.defaultdict(Inventory)
    for post in books_load.iter_postings():
        if post.meta.date >= args.as_of_date:
            continue
        account = post.account
        fund_acct_match = post.account.is_under(*data.FUND_ACCOUNTS)
        is_equity = account.root_part() in data.EQUITY_ACCOUNTS
        if fund_acct_match is None:
            project: MetaValue = None
        else:
            project = post.meta.get(args.meta_key)
            if project is None:
                bc_printer.print_error(Error(
                    post.meta, "no fund specified", post.meta.txn,
                ), file=stderr)
                project = args.unrestricted_fund
            if is_equity:
                if project == args.unrestricted_fund:
                    account = UNRESTRICTED_ACCOUNT
                else:
                    account = RESTRICTED_ACCOUNT
        inventory = inventories[AccountWithFund(account, project)]
        if is_equity:
            inventory.add_amount(post.at_cost())
        else:
            inventory.add_position(get_position(post))

    opening_date = args.as_of_date - datetime.timedelta(1)
    opening = bc_data.Transaction(  # type:ignore[operator]
        None,  # meta
        opening_date,
        '*',
        None,  # payee
        f"Opening balances for FY{fy.for_date(args.as_of_date)}",
        frozenset(),  # tags
        frozenset(),  # links
        [post
         for key in sorted(inventories, key=AccountWithFund.sortkey)
         for post in Posting.build_opening(key, args.meta_key, inventories[key])
        ])
    balance = Balance(get_cost(get_position(post))
                      for post in opening.postings)
    for amount in balance.clean_copy().values():
        opening.postings.append(bc_data.Posting(
            UNRESTRICTED_ACCOUNT, quantize_amount(-amount), None, None, None,
            {args.meta_key: args.unrestricted_fund},
        ))
    dcontext = bc_dcontext.DisplayContext()
    dcontext.set_commas(True)
    bc_printer.print_entry(opening, dcontext, file=stdout)
    return returncode

entry_point = cliutil.make_entry_point(__name__, PROGNAME)

if __name__ == '__main__':
    exit(entry_point())