budget/pfbudget/transform/categorizer.py
Luís Murta c6a13cc83a
Allow multi-category transactions
through the use of tags. Instead of failing when categorizing a
transaction which already has a category, add the new category as a tag
for that transaction.

Issue #2
2024-01-22 22:02:25 +00:00

39 lines
1.2 KiB
Python

from copy import deepcopy
from typing import Iterable, Sequence
from pfbudget.db.model import (
CategoryRule,
CategorySelector,
Transaction,
TransactionCategory,
TransactionTag,
)
from .transform import Transformer
class Categorizer(Transformer):
def __init__(self, rules: Iterable[CategoryRule]):
self.rules = rules
def transform(self, transactions: Sequence[Transaction]) -> Sequence[Transaction]:
result = deepcopy(transactions)
self.transform_inplace(result)
return result
def transform_inplace(self, transactions: Sequence[Transaction]) -> None:
for rule in self.rules:
for transaction in transactions:
if not rule.matches(transaction):
continue
if not transaction.category:
transaction.category = TransactionCategory(
rule.name, CategorySelector.rules
)
else:
if not transaction.tags:
transaction.tags = {TransactionTag(rule.name)}
else:
transaction.tags.add(TransactionTag(rule.name))