62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import re
|
|
|
|
from beancount.core import data
|
|
from beancount.core import inventory
|
|
from beancount.core import account
|
|
from beancount.core import amount
|
|
from beancount.core import account_types
|
|
from beancount.parser import parser
|
|
from beancount.core.number import D
|
|
from beancount.parser import options
|
|
import itertools as it
|
|
from collections import namedtuple
|
|
|
|
LinkAccount = namedtuple("LinkAccount", ["account", "tags"])
|
|
|
|
__plugins__ = ('link_account_plugin',)
|
|
|
|
def add_tags(txn: data.Transaction, *tags: str) -> data.Transaction:
|
|
return txn._replace(tags=txn.tags | set(tags))
|
|
|
|
def link_posting(p: data.Posting, link_accounts: dict[str, LinkAccount]) -> list[data.Posting]:
|
|
if p.account in link_accounts:
|
|
return [
|
|
p._replace(account=link_accounts[p.account].account),
|
|
p._replace(units=-p.units)
|
|
]
|
|
return None
|
|
|
|
|
|
def link_txn(txn: data.Transaction, link_accounts: dict[str, LinkAccount]) -> list[data.Transaction]:
|
|
new_postings = [p for p in txn.postings if p.account in link_accounts]
|
|
if len(new_postings) == 0:
|
|
return []
|
|
|
|
new_txns = []
|
|
|
|
get_account = lambda x: x.account
|
|
new_postings = sorted(new_postings, key=get_account)
|
|
for k, g in it.groupby(new_postings, key=get_account):
|
|
new_postings = [p for o in g for p in link_posting(o, link_accounts)]
|
|
new_txns.append(add_tags(txn._replace(postings=new_postings), "link", *link_accounts[k].tags))
|
|
return new_txns
|
|
|
|
def link_account_plugin(entries, options_map, config=None):
|
|
errors = []
|
|
|
|
link_accounts = {}
|
|
for entry in entries:
|
|
if isinstance(entry, data.Open):
|
|
if "link-account" in entry.meta:
|
|
tags = entry.meta["link-tags"].split(" ") if "link-tags" in entry.meta else []
|
|
link_accounts[entry.account] = LinkAccount(account=entry.meta["link-account"], tags=tags)
|
|
|
|
new_entries = []
|
|
for entry in entries:
|
|
if isinstance(entry, data.Transaction):
|
|
new_entries.extend(link_txn(entry, link_accounts))
|
|
new_entries.append(entry)
|
|
|
|
return new_entries, errors |