262 lines
9.2 KiB
Python
262 lines
9.2 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")
|
|
UnweightedError = collections.namedtuple("UnweightedError", "source message entry")
|
|
SplitAccount = collections.namedtuple("SplitAccount", ["account", "users", "tag"])
|
|
|
|
SPLIT_WEIGHT_CUSTOM = "split-weights"
|
|
|
|
def parse_split_weights(weight_str: str, users: list[str]) -> dict[str, "D"]:
|
|
weights = {}
|
|
for part in weight_str.split(" "):
|
|
user, sep, weight = part.partition(":")
|
|
if not sep:
|
|
raise ValueError(f"Malformed weight entry: {part!r}")
|
|
weights[user] = D(weight)
|
|
if not set(weights) <= set(users):
|
|
raise ValueError(f"Unknown users in weights: {set(weights) - set(users)}")
|
|
return weights
|
|
|
|
|
|
def weights_for_tag_at_date(split_weights, tag, date, users: list[str]):
|
|
history = split_weights.get(tag)
|
|
if not history:
|
|
return None
|
|
applicable = None
|
|
for entry_date, weights in history:
|
|
if entry_date <= date:
|
|
applicable = weights
|
|
else:
|
|
break
|
|
if applicable is None:
|
|
return None
|
|
return [applicable.get(u, D("0")) for u in users]
|
|
|
|
class MultipleAccountException(Exception):
|
|
def __init__(self, split_tags):
|
|
self.split_tags = split_tags
|
|
|
|
class UnweightedException(Exception):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def split_posting(
|
|
posting: data.Posting,
|
|
split: SplitAccount,
|
|
accs: account_types.AccountTypes,
|
|
main: None | str,
|
|
split_weights,
|
|
date,
|
|
) -> list[data.Posting]:
|
|
|
|
weights = weights_for_tag_at_date(split_weights, split.tag, date, split.users)
|
|
if weights is None:
|
|
weights = [1 for _ in range(len(split.users))]
|
|
|
|
if "split" in posting.meta:
|
|
assert(set(posting.meta["split"].split(" ")) <= set(split.users))
|
|
users = posting.meta["split"].split(" ")
|
|
weights = [w if user in users else 0 for (w, user) in zip(weights, 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]
|
|
|
|
total_weight = sum(weights)
|
|
if total_weight == 0:
|
|
raise UnweightedException()
|
|
|
|
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,
|
|
split_weights,
|
|
) -> 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, split_weights, txn.date))
|
|
|
|
# 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)
|
|
|
|
split_weights = collections.defaultdict(list)
|
|
for entry in entries:
|
|
if isinstance(entry, data.Custom) and entry.type == SPLIT_WEIGHT_CUSTOM:
|
|
if len(entry.values) != 2:
|
|
errors.append(ConfigError(
|
|
entry.meta,
|
|
f"Split weight custom entry needs exactly 2 arguments (tag, weights), got {len(entry.values)}",
|
|
entry
|
|
))
|
|
continue
|
|
tag = entry.values[0].value
|
|
weight_str = entry.values[1].value
|
|
if tag not in split_tags:
|
|
errors.append(ConfigError(
|
|
entry.meta,
|
|
f"Split weight custom entry refers to unknown split tag: {tag}",
|
|
entry
|
|
))
|
|
continue
|
|
users = split_accounts[split_tags[tag]].users
|
|
try:
|
|
weights = parse_split_weights(weight_str, users)
|
|
except ValueError as e:
|
|
errors.append(ConfigError(entry.meta, f"Invalid split weights for tag {tag}: {e}", entry))
|
|
continue
|
|
split_weights[tag].append((entry.date, weights))
|
|
|
|
for tag in split_weights:
|
|
split_weights[tag].sort(key=lambda dw: dw[0])
|
|
|
|
new_accounts = set()
|
|
for entry in entries:
|
|
if isinstance(entry, data.Transaction):
|
|
try:
|
|
entry = split_txn(entry, split_accounts, split_tags, accs, config, split_weights)
|
|
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
|
|
))
|
|
except UnweightedException as e:
|
|
errors.append(UnweightedError(
|
|
entry.meta,
|
|
f"Split transaction had no payer as all involved users had no contribution",
|
|
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
|