beancount plugin

non_leaf_account_alias.py

"""add non-leaf accounts a alias

Case:
    * If you open a account `Expenses:Food:Dinner`, and then add a posting with `Expenses:Food`
    * If you open a account `Expenses:Food` and `Expense:Food:Dinner`, and get incomprehensible balance in fava

Example 1:
    2021-10-01 open Assets:Cash
    2021-10-01 open Expenses:Food:Dinner
    2021-10-01 open Expenses:Food:Blah

    2021-10-10 * "I don’t know which category it should be classified into, and I’m too lazy to open an account"
      Asset:Cash -100 USD
      Expenses:Food

    2021-10-11 * "I forgot to assign to the detailed account "
      Asset:Cash -10 USD
      Expenses:Food

Then this plugin will transform the transaction into:

    2021-10-10 open Expenses:Food:Unspecified

    2021-10-10 * "I don’t know which category it should be classified into, and I’m too lazy to open an account"
      Asset:Cash -100 USD
      Expenses:Food:Unspecified

    2021-10-11 * "I forgot to assign to the detailed account "
      Asset:Cash -10 USD
      Expenses:Food:Unspecified

Example 2:
    2020-10-01 open Assets:Cash

    ; simply opening an account
    2020-10-01 open Expenses:Food

    2020-10-01 * "Blah Restaurant" "Dinner"
      Asset:Cash -100 USD
      Expenses:Food

    ; ... (a large number of transactions)

    ; Suddenly want to open a sub-account
    2021-10-01 open Expenses:Food:Blah

Then this plugin will transform the transaction into:

    2020-10-01 open Expenses:Food:Unspecified

    2020-10-01 * "Blah Restaurant" "Dinner"
      Asset:Cash -100 USD
      Expenses:Food:Unspecified

    ; ... (a large number of transactions)

You can just write it down like this. Or use other tools to move `Expenses:Food:Unspecified` to the corresponding other sub-accounts.
"""
from typing import List

from beancount.core import account
from beancount.core import data
from beancount.core import realization

__plugins__ = ('non_leaf_account_alias',)


def load_config(config_string="", default={}):
    if len(config_string) == 0:
        config_obj = {}
    else:
        config_obj = eval(config_string, {}, {})

    if not isinstance(config_obj, dict):
        config_obj = {}

    return config_obj


def get_non_leaf_accounts(real_account: realization.RealAccount) -> List[str]:
    accounts = []
    for key, real_child in real_account.items():
        if len(real_child.items()) > 0:
            accounts.append(real_child.account)
            for real_sub_child_account in get_non_leaf_accounts(real_child):
                accounts.append(real_sub_child_account)

    return accounts


def non_leaf_account_alias(entries: data.Entries, options_map, config_string=""):
    config = load_config(config_string)
    config['account_suffix'] = config.get('account_suffix', 'Unspecified')
    config['auto_open_alias_account'] = config.get('auto_open_alias_account', True)
    new_entries = []
    errors = []
    real_root = realization.realize(entries, compute_balance=False)
    # sys.stderr.write('%s\n' % str(real_root))

    non_leaf_accounts = get_non_leaf_accounts(real_root)
    # sys.stderr.write('%s\n' % str(non_leaf_accounts))

    ready_to_modified = []
    for entry in entries:
        if isinstance(entry, data.Transaction):
            posting: data.Posting
            for i, posting in enumerate(entry.postings):
                if posting.account in non_leaf_accounts:
                    # Using non_leaf accounts.
                    ready_to_modified.append((entry, i))

    if len(ready_to_modified) > 0:
        new_entries.extend(entries)
        opened_accounts = {entry.account
                           for entry in entries
                           if isinstance(entry, data.Open)}
        # start replace non_leaf accounts with '*:Unspecified'
        for entry, i in ready_to_modified:
            popped_posting: data.Posting = entry.postings.pop(i)
            new_account = popped_posting.account
            while new_account in non_leaf_accounts:
                new_account = str.join(account.sep, [new_account, config['account_suffix']])
            if config['auto_open_alias_account'] and new_account not in opened_accounts:
                opened_accounts.add(new_account)
                # open new account
                meta = data.new_metadata(entry.meta['filename'], entry.meta['lineno'])
                new_entries.append(data.Open(meta, entry.date, new_account, None, None))
            entry.postings.insert(i, popped_posting._replace(account=new_account))

        new_entries.sort(key=data.entry_sortkey)
    else:
        new_entries = entries

    return new_entries, errors