from pfbudget.db.model import ( Category, CategorySelector, Selector, Tag, Transaction, TransactionCategory, TransactionTag, ) from codetiming import Timer from datetime import timedelta Transactions = list[Transaction] class Categorizer: options = {} def __init__(self): self.options["null_days"] = 4 def rules( self, transactions: Transactions, categories: list[Category], tags: list[Tag], ): """Overarching categorization tool Receives a list of transactions (by ref) and updates their category according to the rules defined for each category Args: transactions (list[Transaction]): uncategorized transactions categories (list[Category]): available categories tags (list[Tag]): currently available tags """ self._nullify(transactions) self._rule_based_categories(transactions, categories) self._rule_based_tags(transactions, tags) def manual( self, transactions: Transactions, categories: list[Category], tags: list[Tag], ): """Manual categorization input Args: transactions (list[Transaction]): uncategorized transactions categories (list[Category]): available categories tags (list[Tag]): currently available tags """ self._manual(transactions) @Timer(name="nullify") def _nullify(self, transactions: Transactions): count = 0 matching = [] for transaction in transactions: for cancel in ( cancel for cancel in transactions if ( transaction.date - timedelta(days=self.options["null_days"]) <= cancel.date <= transaction.date + timedelta(days=self.options["null_days"]) and transaction not in matching and cancel not in matching and cancel != transaction and cancel.bank != transaction.bank and cancel.amount == -transaction.amount ) ): transaction.category = TransactionCategory( name="null", selector=CategorySelector(Selector.nullifier) ) cancel.category = TransactionCategory( name="null", selector=CategorySelector(Selector.nullifier) ) matching.extend([transaction, cancel]) count += 2 break print(f"Nullified {count} transactions") @Timer(name="categoryrules") def _rule_based_categories( self, transactions: Transactions, categories: list[Category] ): d = {} for category in [c for c in categories if c.rules]: for rule in category.rules: # for transaction in [t for t in transactions if not t.category]: for transaction in [ t for t in transactions if not t.category or t.category.name != "null" ]: if not rule.matches(transaction): continue # passed all conditions, assign category if ( transaction.category and transaction.category.name == category.name ): if ( input(f"Overwrite {transaction} with {category}? (y/n)") == "y" ): transaction.category.name = category.name transaction.category.selector.selector = Selector.rules else: transaction.category = TransactionCategory( category.name, CategorySelector(Selector.rules) ) if rule in d: d[rule] += 1 else: d[rule] = 1 for k, v in d.items(): print(f"{v}: {k}") @Timer(name="tagrules") def _rule_based_tags(self, transactions: Transactions, tags: list[Tag]): d = {} for tag in [t for t in tags if t.rules]: for rule in tag.rules: # for transaction in [t for t in transactions if not t.category]: for transaction in [ t for t in transactions if tag.name not in [tag.tag for tag in t.tags] ]: if not rule.matches(transaction): continue if not transaction.tags: transaction.tags = {TransactionTag(tag.name)} else: transaction.tags.add(TransactionTag(tag.name)) if rule in d: d[rule] += 1 else: d[rule] = 1 for k, v in d.items(): print(f"{v}: {k}") def _manual(self, transactions: Transactions): uncategorized = [t for t in transactions if not t.category] print(f"{len(uncategorized)} transactions left to categorize") for transaction in uncategorized: while True: category = input(f"{transaction} category: ") if category == "quit": return if not category: print("{category} doesn't exist") continue transaction.category = TransactionCategory( category, CategorySelector(Selector.manual) ) break