diff --git a/split_account.py b/split_account.py index 52b05c9..eaf319a 100644 --- a/split_account.py +++ b/split_account.py @@ -19,31 +19,69 @@ __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 + posting: data.Posting, + split: SplitAccount, + accs: account_types.AccountTypes, + main: None | str, + split_weights, + date, ) -> list[data.Posting]: - weights = None + + 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 = [1 if user in users else 0 for user in split.users] + 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] - else: - weights = [1 for _ in range(len(split.users))] total_weight = sum(weights) - assert(total_weight > 0) + if total_weight == 0: + raise UnweightedException() new_postings = [ posting._replace( @@ -93,11 +131,12 @@ def merge_postings(postings: Sequence[data.Posting]) -> list[data.Posting]: def split_txn( - txn: data.Transaction, - split_accounts: dict[str, SplitAccount], - split_tags: dict[str, str], - accs: account_types.AccountTypes, - main: None | str + 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) @@ -116,7 +155,7 @@ def split_txn( split_postings = [] for acc in involved_accs: ps = [(p._replace(account="") 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_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)) @@ -157,11 +196,41 @@ def split_account_plugin(entries, options_map, config=None): ) 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) + 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( @@ -169,6 +238,12 @@ def split_account_plugin(entries, options_map, config=None): 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]