264 lines
9.9 KiB
Python
264 lines
9.9 KiB
Python
"""Learn conservative YNAB category suggestions from transaction CSV history."""
|
|
|
|
import csv
|
|
from dataclasses import dataclass
|
|
from decimal import Decimal, InvalidOperation, ROUND_HALF_UP
|
|
from difflib import SequenceMatcher
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
import re
|
|
import unicodedata
|
|
|
|
|
|
_MEMORY_VERSION = 1
|
|
_MIN_PAYEE_SIMILARITY = 0.70
|
|
_MIN_CATEGORY_SHARE = 0.80
|
|
_AMOUNT_SMOOTHING_CENTS = 100
|
|
|
|
|
|
@dataclass(frozen=True, order=True)
|
|
class CategoryPattern:
|
|
"""An aggregated payee/category/amount pattern, without transaction dates."""
|
|
|
|
payee: str
|
|
category: str
|
|
direction: str
|
|
amount_cents: int
|
|
count: int = 1
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CategorySuggestion:
|
|
"""A category candidate and whether evidence is strong enough to apply it."""
|
|
|
|
category_name: str
|
|
confidence: float
|
|
category_id: str | None
|
|
applied: bool
|
|
|
|
|
|
def normalize_payee(payee: str) -> str:
|
|
"""Normalize formatting and discard long numeric or mixed reference codes."""
|
|
normalized = unicodedata.normalize("NFKD", payee).encode("ascii", "ignore").decode()
|
|
tokens = re.findall(r"[a-z0-9]+", normalized.casefold())
|
|
stable_tokens = [
|
|
token
|
|
for token in tokens
|
|
if not (
|
|
(token.isdigit() and len(token) >= 3)
|
|
or (
|
|
len(token) >= 5
|
|
and re.search(r"[a-z]", token)
|
|
and re.search(r"[0-9]", token)
|
|
)
|
|
)
|
|
]
|
|
return " ".join(stable_tokens)
|
|
|
|
|
|
def _amount_cents(value: str) -> int | None:
|
|
cleaned = value.strip().replace("$", "").replace(",", "")
|
|
if not cleaned:
|
|
return None
|
|
try:
|
|
amount = Decimal(cleaned)
|
|
except InvalidOperation:
|
|
return None
|
|
if not amount.is_finite():
|
|
return None
|
|
cents = (amount * 100).quantize(Decimal("1"), rounding=ROUND_HALF_UP)
|
|
return int(cents)
|
|
|
|
|
|
class CategoryHistory:
|
|
"""Memory of unique merchant/category/amount patterns and snapshot support counts."""
|
|
|
|
def __init__(
|
|
self,
|
|
patterns: dict[tuple[str, str, str, int], CategoryPattern] | None = None,
|
|
) -> None:
|
|
self.patterns = patterns or {}
|
|
|
|
@staticmethod
|
|
def _key(pattern: CategoryPattern) -> tuple[str, str, str, int]:
|
|
return (
|
|
pattern.payee,
|
|
pattern.category.casefold(),
|
|
pattern.direction,
|
|
pattern.amount_cents,
|
|
)
|
|
|
|
@staticmethod
|
|
def memory_path_for_csv(csv_path: str | Path) -> Path:
|
|
"""Return the portable category-memory sidecar path beside the CSV."""
|
|
return Path(csv_path).with_name("category_memory.json")
|
|
|
|
@classmethod
|
|
def load(cls, path: str | Path) -> "CategoryHistory":
|
|
"""Load and validate a category-memory sidecar."""
|
|
memory_path = Path(path)
|
|
payload = json.loads(memory_path.read_text(encoding="utf-8"))
|
|
if not isinstance(payload, dict) or payload.get("version") != _MEMORY_VERSION:
|
|
raise ValueError(f"Unsupported category memory version in {memory_path}")
|
|
|
|
patterns: dict[tuple[str, str, str, int], CategoryPattern] = {}
|
|
for item in payload.get("patterns", []):
|
|
payee = str(item["payee"]).strip()
|
|
category = str(item["category"]).strip()
|
|
direction = str(item["direction"]).strip().lower()
|
|
amount_cents = int(item["amount_cents"])
|
|
count = max(1, int(item.get("count", 1)))
|
|
if payee and category and direction in ("inflow", "outflow") and amount_cents > 0:
|
|
pattern = CategoryPattern(payee, category, direction, amount_cents, count)
|
|
patterns[cls._key(pattern)] = pattern
|
|
return cls(patterns)
|
|
|
|
@classmethod
|
|
def from_csv(cls, csv_path: str | Path) -> "CategoryHistory":
|
|
"""Read the repository's four-column transaction CSV as category examples."""
|
|
path = Path(csv_path)
|
|
patterns: dict[tuple[str, str, str, int], CategoryPattern] = {}
|
|
with path.open("r", encoding="utf-8-sig", newline="") as csv_file:
|
|
reader = csv.DictReader(csv_file)
|
|
required = {"payee", "category", "inflow", "outflow"}
|
|
if not reader.fieldnames or not required.issubset(reader.fieldnames):
|
|
raise ValueError(
|
|
f"{path} must contain payee, category, inflow, and outflow columns"
|
|
)
|
|
|
|
for row in reader:
|
|
payee = normalize_payee(row.get("payee") or "")
|
|
category = (row.get("category") or "").strip()
|
|
inflow = _amount_cents(row.get("inflow") or "")
|
|
outflow = _amount_cents(row.get("outflow") or "")
|
|
if not payee or not category:
|
|
continue
|
|
if inflow is not None and inflow > 0 and (outflow is None or outflow == 0):
|
|
pattern = CategoryPattern(payee, category, "inflow", inflow)
|
|
key = cls._key(pattern)
|
|
existing = patterns.get(key)
|
|
patterns[key] = CategoryPattern(
|
|
payee, category, "inflow", inflow,
|
|
(existing.count if existing else 0) + 1,
|
|
)
|
|
elif outflow is not None and outflow > 0 and (inflow is None or inflow == 0):
|
|
pattern = CategoryPattern(payee, category, "outflow", outflow)
|
|
key = cls._key(pattern)
|
|
existing = patterns.get(key)
|
|
patterns[key] = CategoryPattern(
|
|
payee, category, "outflow", outflow,
|
|
(existing.count if existing else 0) + 1,
|
|
)
|
|
return cls(patterns)
|
|
|
|
def merge(self, other: "CategoryHistory") -> int:
|
|
"""Union patterns, returning how many new patterns were added."""
|
|
changes = 0
|
|
for key, pattern in other.patterns.items():
|
|
existing = self.patterns.get(key)
|
|
if existing is None or pattern.count > existing.count:
|
|
self.patterns[key] = pattern
|
|
changes += 1
|
|
return changes
|
|
|
|
def save(self, path: str | Path) -> None:
|
|
"""Write a deterministic, portable JSON sidecar using an atomic replace."""
|
|
memory_path = Path(path)
|
|
memory_path.parent.mkdir(parents=True, exist_ok=True)
|
|
payload = {
|
|
"version": _MEMORY_VERSION,
|
|
"patterns": [
|
|
{
|
|
"payee": pattern.payee,
|
|
"category": pattern.category,
|
|
"direction": pattern.direction,
|
|
"amount_cents": pattern.amount_cents,
|
|
"count": pattern.count,
|
|
}
|
|
for pattern in sorted(
|
|
self.patterns.values(),
|
|
key=lambda item: (
|
|
item.payee,
|
|
item.category.casefold(),
|
|
item.direction,
|
|
item.amount_cents,
|
|
),
|
|
)
|
|
],
|
|
}
|
|
temporary_path = memory_path.with_name(memory_path.name + ".tmp")
|
|
temporary_path.write_text(
|
|
json.dumps(payload, indent=2, ensure_ascii=False) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
temporary_path.replace(memory_path)
|
|
|
|
def suggest(
|
|
self,
|
|
payee: str,
|
|
amount_milli: int,
|
|
category_ids_by_name: dict[str, list[str]],
|
|
) -> CategorySuggestion | None:
|
|
"""Rank historical categories using payee similarity, direction, and amount."""
|
|
payee_key = normalize_payee(payee)
|
|
if not payee_key or not self.patterns:
|
|
return None
|
|
|
|
direction = "inflow" if amount_milli > 0 else "outflow"
|
|
amount_cents = round(abs(amount_milli) / 10)
|
|
category_scores: dict[str, float] = {}
|
|
category_names: dict[str, str] = {}
|
|
category_support: dict[str, int] = {}
|
|
category_exact: dict[str, bool] = {}
|
|
|
|
for pattern in self.patterns.values():
|
|
if pattern.direction != direction:
|
|
continue
|
|
payee_similarity = SequenceMatcher(
|
|
None, payee_key, pattern.payee, autojunk=False
|
|
).ratio()
|
|
if payee_similarity < _MIN_PAYEE_SIMILARITY:
|
|
continue
|
|
|
|
ratio = (amount_cents + _AMOUNT_SMOOTHING_CENTS) / (
|
|
pattern.amount_cents + _AMOUNT_SMOOTHING_CENTS
|
|
)
|
|
amount_similarity = 1 / (1 + abs(math.log(ratio)))
|
|
weight = payee_similarity * (amount_similarity ** 2) * math.sqrt(pattern.count)
|
|
category_key = pattern.category.casefold()
|
|
category_scores[category_key] = category_scores.get(category_key, 0.0) + weight
|
|
category_names.setdefault(category_key, pattern.category)
|
|
category_support[category_key] = (
|
|
category_support.get(category_key, 0) + pattern.count
|
|
)
|
|
if payee_key == pattern.payee:
|
|
category_exact[category_key] = True
|
|
else:
|
|
category_exact.setdefault(category_key, False)
|
|
|
|
total_score = sum(category_scores.values())
|
|
if total_score <= 0:
|
|
return None
|
|
|
|
best_category = max(category_scores, key=category_scores.get)
|
|
confidence = category_scores[best_category] / total_score
|
|
sufficiently_supported = (
|
|
category_exact.get(best_category, False)
|
|
or category_support[best_category] >= 2
|
|
)
|
|
applied = (
|
|
confidence >= _MIN_CATEGORY_SHARE
|
|
and sufficiently_supported
|
|
)
|
|
|
|
category_name = category_names[best_category]
|
|
matching_ids = category_ids_by_name.get(category_name.casefold(), [])
|
|
category_id = matching_ids[0] if applied and len(matching_ids) == 1 else None
|
|
return CategorySuggestion(
|
|
category_name=category_name,
|
|
confidence=confidence,
|
|
category_id=category_id,
|
|
applied=category_id is not None,
|
|
)
|