from __future__ import annotations from collections.abc import Mapping, MutableMapping, Sequence from dataclasses import dataclass import datetime as dt import decimal import enum import re from typing import Annotated, Any, Callable, Optional, Self, cast from sqlalchemy import ( BigInteger, Enum, ForeignKey, Integer, MetaData, Numeric, String, Text, ) from sqlalchemy.orm import ( DeclarativeBase, Mapped, mapped_column, MappedAsDataclass, relationship, ) class Base(MappedAsDataclass, DeclarativeBase): metadata = MetaData( schema="pfbudget", naming_convention={ "ix": "ix_%(column_0_label)s", "uq": "uq_%(table_name)s_%(column_0_name)s", "ck": "ck_%(table_name)s_`%(constraint_name)s`", "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", "pk": "pk_%(table_name)s", }, ) type_annotation_map = { enum.Enum: Enum(enum.Enum, create_constraint=True, inherit_schema=True), } @dataclass class Serializable: def serialize(self) -> Mapping[str, Any]: return dict(class_=type(self).__name__) @classmethod def deserialize(cls, map: Mapping[str, Any]) -> Self: raise NotImplementedError class AccountType(enum.Enum): checking = enum.auto() savings = enum.auto() investment = enum.auto() mealcard = enum.auto() VISA = enum.auto() MASTERCARD = enum.auto() class Bank(Base, Serializable): __tablename__ = "banks" name: Mapped[str] = mapped_column(primary_key=True) BIC: Mapped[str] = mapped_column(String(8)) type: Mapped[AccountType] nordigen: Mapped[Optional[Nordigen]] = relationship(default=None, lazy="joined") def serialize(self) -> Mapping[str, Any]: nordigen = None if self.nordigen: nordigen = { "bank_id": self.nordigen.bank_id, "requisition_id": self.nordigen.requisition_id, "invert": self.nordigen.invert, } return super().serialize() | dict( name=self.name, BIC=self.BIC, type=self.type.name, nordigen=nordigen, ) @classmethod def deserialize(cls, map: Mapping[str, Any]) -> Self: bank = cls(map["name"], map["BIC"], map["type"]) if map["nordigen"]: bank.nordigen = Nordigen(**map["nordigen"]) return bank bankfk = Annotated[str, mapped_column(Text, ForeignKey(Bank.name))] idpk = Annotated[ int, mapped_column( BigInteger().with_variant(Integer, "sqlite"), primary_key=True, autoincrement=True, ), ] money = Annotated[decimal.Decimal, mapped_column(Numeric(16, 2))] class Transaction(Base, Serializable): __tablename__ = "transactions" id: Mapped[idpk] = mapped_column(init=False) date: Mapped[dt.date] description: Mapped[Optional[str]] amount: Mapped[money] split: Mapped[bool] = mapped_column(default=False) category: Mapped[Optional[TransactionCategory]] = relationship( back_populates="transaction", default=None, lazy="joined" ) tags: Mapped[set[TransactionTag]] = relationship(default_factory=set, lazy="joined") note: Mapped[Optional[Note]] = relationship( cascade="all, delete-orphan", passive_deletes=True, default=None, lazy="joined" ) type: Mapped[str] = mapped_column(init=False) __mapper_args__ = {"polymorphic_on": "type", "polymorphic_identity": "transaction"} def serialize(self) -> Mapping[str, Any]: category = None if self.category: category = { "name": self.category.name, "selector": self.category.selector.name, } return super().serialize() | dict( id=self.id, date=self.date.isoformat(), description=self.description, amount=str(self.amount), split=self.split, category=category if category else None, tags=[{"tag": tag.tag} for tag in self.tags], note={"note": self.note.note} if self.note else None, type=self.type, ) @classmethod def deserialize( cls, map: Mapping[str, Any] ) -> Transaction | BankTransaction | MoneyTransaction | SplitTransaction: match map["type"]: case "bank": return BankTransaction.deserialize(map) case "money": return MoneyTransaction.deserialize(map) case "split": return SplitTransaction.deserialize(map) case _: return cls._deserialize(map) @classmethod def _deserialize(cls, map: Mapping[str, Any]) -> Self: category = None if map["category"]: category = TransactionCategory(map["category"]["name"]) if map["category"]["selector"]: category.selector = map["category"]["selector"] tags: set[TransactionTag] = set() if map["tags"]: tags = set(TransactionTag(t["tag"]) for t in map["tags"]) note = None if map["note"]: note = Note(map["note"]["note"]) result = cls( dt.date.fromisoformat(map["date"]), map["description"], map["amount"], map["split"], category, tags, note, ) if map["id"]: result.id = map["id"] return result def __lt__(self, other: Transaction): return self.date < other.date idfk = Annotated[ int, mapped_column(BigInteger, ForeignKey(Transaction.id, ondelete="CASCADE")) ] class BankTransaction(Transaction): bank: Mapped[Optional[bankfk]] = mapped_column(default=None) __mapper_args__ = {"polymorphic_identity": "bank", "polymorphic_load": "inline"} def serialize(self) -> Mapping[str, Any]: map = cast(MutableMapping[str, Any], super().serialize()) map["bank"] = self.bank return map @classmethod def deserialize(cls, map: Mapping[str, Any]) -> Self: transaction = cls._deserialize(map) transaction.bank = map["bank"] return transaction class MoneyTransaction(Transaction): __mapper_args__ = {"polymorphic_identity": "money"} def serialize(self) -> Mapping[str, Any]: return super().serialize() @classmethod def deserialize(cls, map: Mapping[str, Any]) -> Self: return cls._deserialize(map) class SplitTransaction(Transaction): original: Mapped[Optional[idfk]] = mapped_column(default=None) __mapper_args__ = {"polymorphic_identity": "split", "polymorphic_load": "inline"} def serialize(self) -> Mapping[str, Any]: map = cast(MutableMapping[str, Any], super().serialize()) map["original"] = self.original return map @classmethod def deserialize(cls, map: Mapping[str, Any]) -> Self: transaction = cls._deserialize(map) transaction.original = map["original"] return transaction class CategoryGroup(Base, Serializable): __tablename__ = "category_groups" name: Mapped[str] = mapped_column(primary_key=True) def serialize(self) -> Mapping[str, Any]: return super().serialize() | dict(name=self.name) @classmethod def deserialize(cls, map: Mapping[str, Any]) -> Self: return cls(map["name"]) class Category(Base, Serializable, repr=False): __tablename__ = "categories" name: Mapped[str] = mapped_column(primary_key=True) group: Mapped[Optional[str]] = mapped_column( ForeignKey(CategoryGroup.name), default=None ) rules: Mapped[list[CategoryRule]] = relationship( cascade="all, delete-orphan", passive_deletes=True, default_factory=list, lazy="joined", ) schedule: Mapped[Optional[CategorySchedule]] = relationship( cascade="all, delete-orphan", passive_deletes=True, default=None, lazy="joined" ) def serialize(self) -> Mapping[str, Any]: rules: Sequence[Mapping[str, Any]] = [] for rule in self.rules: rules.append( { "start": rule.start.isoformat() if rule.start else None, "end": rule.end.isoformat() if rule.end else None, "description": rule.description, "regex": rule.regex, "bank": rule.bank, "min": str(rule.min) if rule.min is not None else None, "max": str(rule.max) if rule.max is not None else None, } ) schedule = None if self.schedule: schedule = { "period": self.schedule.period.name if self.schedule.period else None, "period_multiplier": self.schedule.period_multiplier, "amount": self.schedule.amount, } return super().serialize() | dict( name=self.name, group=self.group, rules=rules, schedule=schedule, ) @classmethod def deserialize(cls, map: Mapping[str, Any]) -> Self: rules: list[CategoryRule] = [] for rule in map["rules"]: rules.append( CategoryRule( dt.date.fromisoformat(rule["start"]) if rule["start"] else None, dt.date.fromisoformat(rule["end"]) if rule["end"] else None, rule["description"], rule["regex"], rule["bank"], rule["min"], rule["max"], ) ) return cls( map["name"], map["group"], rules, CategorySchedule(**map["schedule"]) if map["schedule"] else None, ) def __repr__(self) -> str: return ( f"Category(name={self.name}, group={self.group}, #rules={len(self.rules)}," f" schedule={self.schedule})" ) catfk = Annotated[ str, mapped_column(ForeignKey(Category.name, ondelete="CASCADE")), ] class CategorySelector(enum.Enum): unknown = enum.auto() nullifier = enum.auto() vacations = enum.auto() rules = enum.auto() algorithm = enum.auto() manual = enum.auto() class TransactionCategory(Base): __tablename__ = "transactions_categorized" id: Mapped[idfk] = mapped_column(primary_key=True, init=False) name: Mapped[catfk] selector: Mapped[CategorySelector] = mapped_column(default=CategorySelector.unknown) transaction: Mapped[Transaction] = relationship( back_populates="category", init=False, compare=False ) class Note(Base): __tablename__ = "notes" id: Mapped[idfk] = mapped_column(primary_key=True, init=False) note: Mapped[str] class Nordigen(Base): __tablename__ = "banks_nordigen" name: Mapped[bankfk] = mapped_column(primary_key=True, init=False) bank_id: Mapped[Optional[str]] requisition_id: Mapped[Optional[str]] invert: Mapped[Optional[bool]] = mapped_column(default=None) class Tag(Base, Serializable): __tablename__ = "tags" name: Mapped[str] = mapped_column(primary_key=True) rules: Mapped[list[TagRule]] = relationship( cascade="all, delete-orphan", passive_deletes=True, default_factory=list, lazy="joined", ) def serialize(self) -> Mapping[str, Any]: rules: Sequence[Mapping[str, Any]] = [] for rule in self.rules: rules.append( { "start": rule.start, "end": rule.end, "description": rule.description, "regex": rule.regex, "bank": rule.bank, "min": str(rule.min) if rule.min is not None else None, "max": str(rule.max) if rule.max is not None else None, } ) return super().serialize() | dict(name=self.name, rules=rules) @classmethod def deserialize(cls, map: Mapping[str, Any]) -> Self: rules: list[TagRule] = [] for rule in map["rules"]: rules.append( TagRule( dt.date.fromisoformat(rule["start"]) if rule["start"] else None, dt.date.fromisoformat(rule["end"]) if rule["end"] else None, rule["description"], rule["regex"], rule["bank"], rule["min"], rule["max"], ) ) return cls(map["name"], rules) class TransactionTag(Base, unsafe_hash=True): __tablename__ = "transactions_tagged" id: Mapped[idfk] = mapped_column(primary_key=True, init=False) tag: Mapped[str] = mapped_column(ForeignKey(Tag.name), primary_key=True) class SchedulePeriod(enum.Enum): daily = enum.auto() weekly = enum.auto() monthly = enum.auto() yearly = enum.auto() class CategorySchedule(Base): __tablename__ = "category_schedules" name: Mapped[catfk] = mapped_column(primary_key=True, init=False) period: Mapped[Optional[SchedulePeriod]] period_multiplier: Mapped[Optional[int]] amount: Mapped[Optional[int]] class Link(Base): __tablename__ = "links" original: Mapped[idfk] = mapped_column(primary_key=True) link: Mapped[idfk] = mapped_column(primary_key=True) class Rule(Base): __tablename__ = "rules" id: Mapped[idpk] = mapped_column(init=False) start: Mapped[Optional[dt.date]] = mapped_column(default=None) end: Mapped[Optional[dt.date]] = mapped_column(default=None) description: Mapped[Optional[str]] = mapped_column(default=None) regex: Mapped[Optional[str]] = mapped_column(default=None) bank: Mapped[Optional[str]] = mapped_column(default=None) min: Mapped[Optional[money]] = mapped_column(default=None) max: Mapped[Optional[money]] = mapped_column(default=None) type: Mapped[str] = mapped_column(init=False) __mapper_args__ = { "polymorphic_identity": "rule", "polymorphic_on": "type", } def matches(self, t: BankTransaction) -> bool: valid = None if self.regex: valid = re.compile(self.regex, re.IGNORECASE) ops = ( Rule.exists(self.start, lambda r: r < t.date), Rule.exists(self.end, lambda r: r > t.date), Rule.exists(self.description, lambda r: r == t.description), Rule.exists( valid, lambda r: r.search(t.description) if t.description else False, ), Rule.exists(self.bank, lambda r: r == t.bank), Rule.exists(self.min, lambda r: r < t.amount), Rule.exists(self.max, lambda r: r > t.amount), ) if all(ops): return True return False @staticmethod def exists(r: Optional[Any], op: Callable[[Any], bool]) -> bool: return op(r) if r is not None else True class CategoryRule(Rule): __tablename__ = "category_rules" id: Mapped[int] = mapped_column( BigInteger, ForeignKey(Rule.id, ondelete="CASCADE"), primary_key=True, init=False, ) name: Mapped[catfk] = mapped_column(init=False) __mapper_args__ = { "polymorphic_identity": "category_rule", "polymorphic_load": "selectin", } class TagRule(Rule): __tablename__ = "tag_rules" id: Mapped[int] = mapped_column( BigInteger, ForeignKey(Rule.id, ondelete="CASCADE"), primary_key=True, init=False, ) tag: Mapped[str] = mapped_column( ForeignKey(Tag.name, ondelete="CASCADE"), init=False ) __mapper_args__ = { "polymorphic_identity": "tag_rule", "polymorphic_load": "selectin", }