187 lines
6.5 KiB
Python
187 lines
6.5 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import re
|
|
import collections
|
|
from typing import Sequence
|
|
|
|
from beancount.core import data
|
|
from beancount.core import inventory
|
|
from beancount.core import account
|
|
from beancount.core import amount
|
|
from beancount.core import getters
|
|
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
|
|
|
|
__plugins__ = ('split_account_plugin',)
|
|
|
|
ConfigError = collections.namedtuple("ConfigError", "source message entry")
|
|
MultipleAccountError = collections.namedtuple("MultipleAccountError", "source message entry")
|
|
SplitAccount = collections.namedtuple("SplitAccount", ["account", "users", "tag"])
|
|
|
|
class MultipleAccountException(Exception):
|
|
def __init__(self, split_tags):
|
|
self.split_tags = split_tags
|
|
|
|
def split_posting(
|
|
posting: data.Posting,
|
|
split: SplitAccount,
|
|
accs: account_types.AccountTypes,
|
|
main: None | str
|
|
) -> list[data.Posting]:
|
|
weights = None
|
|
if "split" in posting.meta:
|
|
assert(set(posting.meta["split"].split(" ")) <= set(split.users))
|
|
users = posting.meta["split"].split(" ")
|
|
weights = [1 if user in users else 0 for user in split.users]
|
|
elif account.leaf(posting.account) in split.users:
|
|
user = account.leaf(posting.account)
|
|
weights = [1 if u == user else 0 for u in split.users]
|
|
else:
|
|
weights = [1 for _ in range(len(split.users))]
|
|
|
|
total_weight = sum(weights)
|
|
assert(total_weight > 0)
|
|
|
|
new_postings = [
|
|
posting._replace(
|
|
units=amount.mul(posting.units, -D(w) / total_weight),
|
|
account=account.join(split.account, split.users[idx]),
|
|
)
|
|
for idx, w in enumerate(weights)
|
|
]
|
|
if main in split.users:
|
|
main_idx = split.users.index(main)
|
|
new_postings.extend([
|
|
posting._replace(
|
|
units=amount.mul(posting.units, -D(w) / total_weight),
|
|
)
|
|
for idx, w in enumerate(weights)
|
|
if idx != main_idx
|
|
])
|
|
new_postings.extend([
|
|
posting._replace(
|
|
units=amount.mul(posting.units, D(w) / total_weight),
|
|
account=account.join("Passiu", "Split", split.users[idx]),
|
|
)
|
|
for idx, w in enumerate(weights)
|
|
if idx != main_idx
|
|
])
|
|
|
|
return new_postings
|
|
|
|
def merge_postings(postings: Sequence[data.Posting]) -> list[data.Posting]:
|
|
merged = []
|
|
keyfunc = lambda p: p.account
|
|
for acc, g in it.groupby(sorted(postings, key=keyfunc), key=keyfunc):
|
|
inv = inventory.Inventory()
|
|
for p in g:
|
|
inv.add_position(p)
|
|
|
|
for pos in inv.get_positions():
|
|
merged.append(data.Posting(
|
|
account=acc,
|
|
units=pos.units,
|
|
cost=pos.cost,
|
|
price=None,
|
|
flag=None,
|
|
meta=None
|
|
))
|
|
return merged
|
|
|
|
|
|
def split_txn(
|
|
txn: data.Transaction,
|
|
split_accounts: dict[str, SplitAccount],
|
|
split_tags: dict[str, str],
|
|
accs: account_types.AccountTypes,
|
|
main: None | str
|
|
) -> data.Transaction:
|
|
involved_accs = set(p.account for p in txn.postings if p.account in split_accounts)
|
|
|
|
tags = list(tag for tag in txn.tags if tag in split_tags)
|
|
split_accs = list(split_tags[tag] for tag in tags)
|
|
if len(split_accs) > 1:
|
|
raise MultipleAccountException(tags)
|
|
elif len(split_accs) == 1:
|
|
involved_accs.add(split_accs[0])
|
|
if len(involved_accs) == 0:
|
|
return txn
|
|
|
|
# filter postings that involve parent accounts, this way we can preserve all asset buckets
|
|
new_postings = [p for p in txn.postings if p.account not in involved_accs]
|
|
|
|
split_postings = []
|
|
for acc in involved_accs:
|
|
ps = [(p._replace(account="<invalid>") if p.account in involved_accs else p) for p in txn.postings if p.account != acc]
|
|
split_postings.extend(n for p in ps for n in split_posting(p, split_accounts[acc], accs, main))
|
|
|
|
# merge split_postings grouped by account
|
|
new_postings.extend(merge_postings(split_postings))
|
|
|
|
return txn._replace(postings=new_postings)
|
|
|
|
def split_account_plugin(entries, options_map, config=None):
|
|
errors = []
|
|
new_entries = []
|
|
accs = options.get_account_types(options_map)
|
|
|
|
split_accounts = {}
|
|
split_tags = {}
|
|
for entry in entries:
|
|
if isinstance(entry, data.Open):
|
|
if "split" in entry.meta:
|
|
tag, *users = entry.meta["split"].split(" ")
|
|
users = list(set(users)) # remove repeats
|
|
if len(users) < 2:
|
|
errors.append(ConfigError(
|
|
entry.meta,
|
|
f"Split account need at least two users, only {len(users)} found: {users}",
|
|
entry
|
|
))
|
|
continue
|
|
|
|
split_accounts[entry.account] = SplitAccount(users=users, account=entry.account, tag=tag)
|
|
split_tags[tag] = entry.account
|
|
for user in users:
|
|
new_entries.append(
|
|
data.Open(
|
|
meta=data.new_metadata("<split-account>", 0),
|
|
date=entry.date,
|
|
account=account.join(entry.account, user),
|
|
currencies=entry.currencies,
|
|
booking=entry.booking
|
|
)
|
|
)
|
|
old_accounts = getters.get_accounts(new_entries) | getters.get_accounts(entries)
|
|
|
|
new_accounts = set()
|
|
for entry in entries:
|
|
if isinstance(entry, data.Transaction):
|
|
try:
|
|
entry = split_txn(entry, split_accounts, split_tags, accs, config)
|
|
new_accounts |= set((p.account for p in entry.postings))
|
|
except MultipleAccountException as e:
|
|
errors.append(MultipleAccountError(
|
|
entry.meta,
|
|
f"Transaction has multiple split tags: {e.split_tags}",
|
|
entry
|
|
))
|
|
new_entries.append(entry)
|
|
|
|
oldest_date = getters.get_min_max_dates(new_entries)[0]
|
|
for acc in new_accounts - old_accounts:
|
|
new_entries.append(
|
|
data.Open(
|
|
meta=data.new_metadata("<split-account>", 1),
|
|
date=oldest_date,
|
|
account=acc,
|
|
currencies=None,
|
|
booking=None
|
|
)
|
|
)
|
|
|
|
return new_entries, errors
|