87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import re
|
|
|
|
from beancount.core import data
|
|
from beancount.core import getters
|
|
from beancount.core import amount
|
|
from beancount.core import account_types
|
|
from beancount.core import realization
|
|
from beancount.core.number import D
|
|
from beancount.parser import options
|
|
|
|
__plugins__ = ('split_account_plugin',)
|
|
|
|
def get_group(tag, _pattern=re.compile(r"split-(.*)")):
|
|
m = _pattern.match(tag)
|
|
if m:
|
|
return m.group(1)
|
|
return ""
|
|
|
|
def is_splitable(posting: data.Posting, accs: account_types.AccountTypes) -> bool:
|
|
if posting.meta.get("split", "") == "no":
|
|
return False
|
|
return account_types.is_income_statement_account(posting.account, accs)
|
|
|
|
def add_tag(txn: data.Transaction, tag: str) -> data.Transaction:
|
|
return txn._replace(tags=txn.tags | set([tag]))
|
|
|
|
def divide_and_flip(posting: data.Posting, n: int = 2) -> data.Posting:
|
|
"""Return a new posting with n-th the amount and opposite sign."""
|
|
assert(n>0)
|
|
old_amount = posting.units
|
|
new_amount = amount.div(old_amount, D(-n))
|
|
return posting._replace(units=new_amount)
|
|
|
|
def split_txn(txn: data.Transaction, split_accounts: dict, accs: account_types.AccountTypes) -> None | data.Transaction:
|
|
# involves_account = any((acc in split_accounts.items() for acc in getters.get_accounts([entry])))
|
|
|
|
groups = [g for g in ( get_group(tag) for tag in txn.tags ) if g in split_accounts]
|
|
if len(groups) != 1:
|
|
return None
|
|
group = groups[0]
|
|
|
|
has_ignore_meta = txn.meta.get("split", "") == "no"
|
|
if has_ignore_meta:
|
|
return None
|
|
|
|
# gather all splitable expenses/income, and get their total balance
|
|
new_postings = [posting for posting in txn.postings if is_splitable(posting, accs)]
|
|
if len(new_postings) == 0:
|
|
return None
|
|
|
|
new_balance = realization.compute_postings_balance(new_postings)
|
|
for position in new_balance:
|
|
new_postings.append(data.Posting(
|
|
account=split_accounts[group],
|
|
units=-position.units,
|
|
cost=position.cost,
|
|
price=None,
|
|
flag=None,
|
|
meta=None
|
|
))
|
|
|
|
new_postings = [ divide_and_flip(p) for p in new_postings ]
|
|
|
|
return add_tag(txn._replace(postings=new_postings), "retorn")
|
|
|
|
def split_account_plugin(entries, options_map):
|
|
errors = []
|
|
new_entries = []
|
|
accs = options.get_account_types(options_map)
|
|
|
|
split_accounts = {}
|
|
for entry in entries:
|
|
if isinstance(entry, data.Open):
|
|
if "split" in entry.meta:
|
|
split_accounts[entry.meta["split"]] = entry.account
|
|
|
|
for entry in entries:
|
|
if isinstance(entry, data.Transaction):
|
|
new_txn = split_txn(entry, split_accounts, accs)
|
|
if new_txn is not None:
|
|
new_entries.append(new_txn)
|
|
new_entries.append(entry)
|
|
|
|
return new_entries, errors
|