beancount-forecast.py

"""Adding a forecasting and amortization feature to Beancount via a plugin

see https://github.com/beancount/beancount/blob/ab3fdc613fd408e5f6d8039b2fe7eb37c0b31a5e/experiments/plugins/forecast.py
also https://github.com/beancount/beancount/blob/d841487ccdda04c159de86b1186e7c2ea997a3e2/beancount/parser/lexer.l#L127-L129

This entry filter plugin uses existing syntax to define and automatically
insert future transactions based on a convention.

A User can create a transaction like this:

    plugin "forecast" "{'years':5}"
    ;plugin "amortize"

    2021-01-01 open Assets:Checking USD
        statement-close 12

    2021-08-25 # "Rent"
        frequency: "monthly" ; required
        interval: 1
        Assets:Checking         500 USD
        Income:Other

    2020-01-30 % "Mortgage Payment"
        frequency: "monthly"            ; required
        principal: 500000               ; required
        period: 12 * 25                 ; required
        rate: 0.025 / 12                ; required
        account: "Income:Payment"       ; required
        interest: "Expenses:Interest"   ; required
        balance: "Liabilities:Loan"   ; required

This transaction will be filtered out and replaced with new transactions
for the period specified in the configuration string (default: 5 years)

see https://dateutil.readthedocs.io/en/stable/rrule.html#dateutil.rrule.rrule for metadata

"""

import calendar
from dateutil.rrule import rrule, FREQNAMES
from datetime import date

from amortization.schedule import amortization_schedule

from beancount.core import data
from beancount.core.amount import Amount #, mul
from beancount.core.number import D, ZERO, Decimal
# from beancount import loader

# __plugins__ = ('forecast', ) # the comma is required for this tuple, else it will split string and look for "f"
__plugins__ = ('forecast','amortize')
# __plugins__ = loader.combine_plugins('forecast', 'amortize')

multiplier = {
    'YEARLY': 1,
    'MONTHLY': 12,
    'WEEKLY': 56,
    'DAILY': 365.25,
}

def forecast(entries, options_map, config_str=None):
    """A filter that piggybacks on top of the Beancount input syntax to
    insert forecast entries automatically. This function accepts the return
    value of beancount.loader.load_file() and must return the same type of output.

    Args:
        entries: a list of entry instances
        options_map: a dict of options parsed from the file
        config_str: a dict of plugin-specific options

    Returns:
        A tuple of entries and errors

    """
    errors = []
    config = eval(config_str, {}, {}) if config_str else {}

    # Filter out forecast entries from the list of valid entries
    forecast_entries = []
    filtered_entries = []
    accounts = {} # accounts with statement closing dates

    for entry in entries:
        if (isinstance(entry, data.Open) and entry.meta.get('statement-close')):
            accounts[entry.account] = entry.meta.get('statement-close')

        outlist = (forecast_entries
                    if (isinstance(entry, data.Transaction) and entry.flag == '#')
                    else filtered_entries)
        outlist.append(entry)

    # Generate forecast entries until meta.until
    new_entries = []

    for entry in forecast_entries:
        # Parse the periodicity
        if not 'frequency' in entry.meta:
            new_entries.append(entry)
            continue
        else:
            frequency = entry.meta['frequency'].upper() # TODO: error handling, default to YEARLY?
            freq = FREQNAMES.index(frequency)
            interval = entry.meta.get('interval', 1)
            _count = multiplier.get(frequency) * config.get('years', 5) / interval
            count = entry.meta.get('count', _count)
            until = date.fromisoformat(entry.meta.get('until')) if entry.meta.get('until') else None

            # Generate a new entry for each forecast date
            # TODO: AFTER the last cleared entry
            try:
                forecast_dates = [dt.date() for dt in rrule(freq=freq, dtstart=entry.date, count=count, interval=interval, until=until)]

                for forecast_date in forecast_dates:
                    forecast_entry = entry._replace(date=forecast_date)

                    # add tags
                    for posting in entry.postings:
                        statement_close = accounts.get(posting.account)

                        if(statement_close):
                            month = forecast_date.month
                            link = ""

                            if (forecast_date.day > statement_close):
                                link = calendar.month_name[month + 1 if month < 12 else 1].lower() + "-" + str(forecast_date.year if month < 12 else forecast_date.year + 1)
                            else:
                                link = calendar.month_name[month].lower() + "-" + str(forecast_date.year)

                            new_links = set(entry.links)
                            new_links.add(link)
                            forecast_entry = forecast_entry._replace(links=new_links)

                    # forecast_entry = entry._replace(date=forecast_date, links=new_links)
                    new_entries.append(forecast_entry)

            except:
                print("Error in Transaction Metadata:")
                # print(entry)
                print(entry.meta)
                # print("=====\n")

    # print(forecast_dates)
    # print(new_entries)
    return filtered_entries + new_entries, errors
    # return filtered_entries + new_entries, []




def amortize(entries, options_map, config_str=None):
    """A filter that piggybacks on top of the Beancount input syntax to
    insert amortization entries automatically. This function accepts the return
    value of beancount.loader.load_file() and must return the same type of output.

    Args:
        entries: a list of entry instances
        options_map: a dict of options parsed from the file
        config_str: a dict of plugin-specific options

    Returns:
        A tuple of entries and errors

    """

    # Filter out amortization entries from the list of valid entries
    amortization_entries = []
    filtered_entries = []

    for entry in entries:
        outlist = (amortization_entries
                    if (isinstance(entry, data.Transaction) and entry.flag == '%')
                    else filtered_entries)
        outlist.append(entry)

    # Generate amortization entries
    new_entries = []

    for entry in amortization_entries:
        if not 'rate' in entry.meta:
            new_entries.append(entry)
            continue
        else:
            frequency = entry.meta['frequency'].upper() # TODO: error handling, default to MONTHLY?
            freq = FREQNAMES.index(frequency)
            count = int(entry.meta['period'])
            forecast_dates = [dt.date() for dt in rrule(freq=freq, dtstart=entry.date)]
            original_postings = entry.postings
            print("=====")
            print(entry)
            # print(entry.meta)
            # print(frequency, freq, count)
            # print(forecast_dates)

            # Generate a new entry for each period date
            for number, amount, interest, principal, balance in amortization_schedule(entry.meta['principal'], entry.meta['rate'], count):
                # print(number, amount, interest, principal, balance, forecast_dates[number-1])
                forecast_entry = entry._replace(
                    date=forecast_dates[number-1],
                    narration=f"{number}/{count}",
                    payee=entry.narration,
                    postings=
                        original_postings +
                        [
                            data.Posting(
                                account=entry.meta['account'],
                                units=Amount(number=amount, currency="USD"), # mul(amount, D(-1))
                                cost=None,
                                price=None,
                                flag=None,
                                meta=None,
                            ),
                            data.Posting(
                                account=entry.meta['balance'],
                                units=Amount(number=principal, currency="USD"),
                                cost=None,
                                price=None,
                                flag=None,
                                meta=None,
                            ),
                            data.Posting(
                                account=entry.meta['interest'],
                                units=Amount(number=interest, currency="USD"),
                                cost=None,
                                price=None,
                                flag=None,
                                meta=None,
                            )
                        ]
                )
                new_entries.append(forecast_entry)
                # TODO: new_entries.append balance

            # add Postings:
            # data.create_simple_posting(forecast_entry, "Income:Payment", "10", "USD")
            # Payment
            # Principal
            # Interest

    # print(amortization_entries)
    # print(new_entries)
    # return entries, []
    return filtered_entries + new_entries, []