generated from felixm/defaultpy
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
from src.models import Transaction, UNKNOWN_CATEGORY
|
|
from src.fzf import iterfzf
|
|
from typing import List
|
|
|
|
|
|
def get_categories(transactions: List[Transaction]) -> List[str]:
|
|
categories = set([t.account2 for t in transactions])
|
|
categories.add(UNKNOWN_CATEGORY)
|
|
return list(categories)
|
|
|
|
|
|
def get_sort_categories():
|
|
def sort_categories(row: str, categories: List[str]):
|
|
if learn is None:
|
|
return
|
|
_, _, probs = learn.predict(row)
|
|
cat_to_prob = dict(zip(learn.dls.vocab[1],probs.tolist()))
|
|
categories.sort(key=lambda c: cat_to_prob[c] if c in cat_to_prob else 0.0, reverse=True)
|
|
|
|
learn = None
|
|
try:
|
|
from fastai.text.all import load_learner
|
|
learn = load_learner("export.pkl")
|
|
except ModuleNotFoundError:
|
|
pass
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
return sort_categories
|
|
|
|
|
|
def add_account2(transactions: List[Transaction]):
|
|
categories = get_categories(transactions)
|
|
unmapped_transactions = filter(lambda t: t.account2 == UNKNOWN_CATEGORY, transactions)
|
|
sort_categories = get_sort_categories()
|
|
for t in unmapped_transactions:
|
|
sort_categories(t.row, categories)
|
|
add_account2_interactive(t, categories)
|
|
|
|
|
|
def add_account2_interactive(transaction: Transaction, categories: List[str]):
|
|
t = transaction
|
|
account2 = None
|
|
prompt = f"{t.account1} {t.date} {t.description} {t.debit} > "
|
|
while account2 is None:
|
|
account2 = iterfzf(categories, prompt=prompt)
|
|
transaction.account2 = account2
|
|
print(f"Assigned category '{account2}'.")
|