129 lines
6.8 KiB
Python
129 lines
6.8 KiB
Python
from __future__ import annotations
|
||
|
||
import re
|
||
|
||
from app.domain import Category
|
||
|
||
|
||
CATEGORY_META: dict[Category, tuple[str, str]] = {
|
||
Category.PRODUCTS: ("🛒", "Продукты"),
|
||
Category.ENTERTAINMENT: ("🎉", "Развлечения"),
|
||
Category.CAR: ("🚗", "Машина"),
|
||
Category.CREDITS: ("🏦", "Кредиты"),
|
||
Category.RENOVATION: ("🛠", "Ремонт"),
|
||
Category.OTHER: ("📦", "Другое"),
|
||
Category.INCOME: ("💚", "Доход"),
|
||
}
|
||
|
||
# Сначала проверяются более специфичные слова. Список легко расширять без миграции БД.
|
||
CATEGORY_KEYWORDS: dict[Category, tuple[str, ...]] = {
|
||
Category.PRODUCTS: (
|
||
"ярче", "мария ра", "пятерочка", "пятёрочка", "магнит", "лента",
|
||
"ашан", "metro", "продукт", "еда", "молоко", "хлеб", "мясо",
|
||
"овощ", "фрукт", "кофе", "кафе", "ресторан", "доставка еды",
|
||
"алкоголь", "пиво", "вино", "супермаркет",
|
||
),
|
||
Category.ENTERTAINMENT: (
|
||
"кино", "театр", "концерт", "бар", "клуб", "игра", "steam",
|
||
"подписка", "музей", "развлеч", "боулинг", "караоке", "хобби",
|
||
),
|
||
Category.CAR: (
|
||
"бензин", "топливо", "азс", "газпромнефть", "лукойл", "g-drive",
|
||
"шиномонтаж", "автосервис", "мойка", "запчаст", "масло", "осаго",
|
||
"парковка", "штраф гибдд", "машин", "авто",
|
||
),
|
||
Category.CREDITS: (
|
||
"кредит", "ипотек", "долг", "рассроч", "процент банку", "заём", "займ",
|
||
),
|
||
Category.RENOVATION: (
|
||
"ремонт квартир", "стройматериал", "обои", "краска", "ламинат",
|
||
"сантехник", "электрик", "мебель", "инструмент", "стройка", "леруа", "ремонт",
|
||
),
|
||
}
|
||
|
||
|
||
def normalize(text: str) -> str:
|
||
text = text.casefold().replace("ё", "е")
|
||
return re.sub(r"[^a-zа-я0-9]+", " ", text).strip()
|
||
|
||
|
||
def classify(description: str) -> Category:
|
||
normalized_text = normalize(description)
|
||
normalized = f" {normalized_text} "
|
||
words = normalized_text.split()
|
||
for category, keywords in CATEGORY_KEYWORDS.items():
|
||
for keyword in keywords:
|
||
normalized_keyword = normalize(keyword)
|
||
is_phrase = " " in normalized_keyword
|
||
exact_phrase = is_phrase and f" {normalized_keyword} " in normalized
|
||
word_match = not is_phrase and any(
|
||
word == normalized_keyword
|
||
or (len(normalized_keyword) >= 4 and word.startswith(normalized_keyword))
|
||
for word in words
|
||
)
|
||
if exact_phrase or word_match:
|
||
return category
|
||
return Category.OTHER
|
||
|
||
|
||
def category_label(category: Category) -> str:
|
||
icon, name = CATEGORY_META[category]
|
||
return f"{icon} {name}"
|
||
|
||
|
||
EXPENSE_CATEGORIES = tuple(category for category in Category if category is not Category.INCOME)
|
||
|
||
|
||
# system_key, icon, visible name, keywords. These are seeded for every family budget.
|
||
DEFAULT_SUBCATEGORIES: dict[Category, tuple[tuple[str, str, str, tuple[str, ...]], ...]] = {
|
||
Category.PRODUCTS: (
|
||
("products_food", "🥦", "Еда", ("продукт", "еда", "ярче", "магнит", "пятер", "лента", "молоко", "хлеб", "мясо", "овощ", "фрукт")),
|
||
("products_alcohol", "🍷", "Алкоголь", ("алкоголь", "пиво", "вино", "водка")),
|
||
("products_cafe", "☕", "Кафе и рестораны", ("кафе", "ресторан", "кофе", "доставка еды")),
|
||
("products_household", "🧻", "Бытовые товары", ("бытовая химия", "хозтовар", "порошок", "салфетк")),
|
||
),
|
||
Category.ENTERTAINMENT: (
|
||
("entertainment_cinema", "🎬", "Кино и театр", ("кино", "театр", "концерт", "музей")),
|
||
("entertainment_games", "🎮", "Игры", ("игра", "steam", "playstation", "xbox")),
|
||
("entertainment_subscriptions", "📺", "Подписки", ("подписка", "кинопоиск", "ivi", "netflix")),
|
||
("entertainment_hobby", "🎨", "Хобби", ("хобби", "боулинг", "караоке")),
|
||
),
|
||
Category.CAR: (
|
||
("car_fuel", "⛽", "Бензин", ("бензин", "топливо", "азс", "газпромнефть", "лукойл", "g drive")),
|
||
("car_repair", "🔧", "Ремонт", ("автосервис", "ремонт машины", "ремонт авто", "шиномонтаж")),
|
||
("car_parts", "⚙️", "Запчасти", ("запчаст", "масло", "аккумулятор", "шина", "резина")),
|
||
("car_service", "🅿️", "Мойка и парковка", ("мойка", "парковка")),
|
||
("car_insurance", "🛡", "Страховка", ("осаго", "каско", "страховка авто")),
|
||
),
|
||
Category.CREDITS: (
|
||
("credits_mortgage", "🏠", "Ипотека", ("ипотек",)),
|
||
("credits_bank", "💳", "Кредиты", ("кредит", "процент банку")),
|
||
("credits_installment", "📆", "Рассрочки", ("рассроч",)),
|
||
),
|
||
Category.RENOVATION: (
|
||
("renovation_materials", "🧱", "Материалы", ("стройматериал", "обои", "краска", "ламинат", "леруа")),
|
||
("renovation_furniture", "🪑", "Мебель", ("мебель", "диван", "шкаф", "стол")),
|
||
("renovation_work", "👷", "Работы", ("сантехник", "электрик", "ремонт квартир", "стройка")),
|
||
("renovation_tools", "🧰", "Инструменты", ("инструмент", "дрель", "шуруповерт")),
|
||
),
|
||
Category.OTHER: (),
|
||
}
|
||
|
||
|
||
def suggest_subcategory(category: Category, description: str) -> str | None:
|
||
normalized_text = normalize(description)
|
||
words = normalized_text.split()
|
||
padded = f" {normalized_text} "
|
||
for system_key, _icon, _name, keywords in DEFAULT_SUBCATEGORIES.get(category, ()):
|
||
for keyword in keywords:
|
||
normalized_keyword = normalize(keyword)
|
||
if " " in normalized_keyword:
|
||
if f" {normalized_keyword} " in padded:
|
||
return system_key
|
||
elif any(
|
||
word == normalized_keyword
|
||
or (len(normalized_keyword) >= 4 and word.startswith(normalized_keyword))
|
||
for word in words
|
||
):
|
||
return system_key
|
||
return None
|