homebank/src/homebank/Homebank.py

from __future__ import annotations

import datetime
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from lxml import etree


# XML functions
def get_attribute(node: etree.Element, name: str) -> str:
    if not name in node.attrib:
        raise KeyError(f"Node '{node.tag}' in line {node.sourceline} is missing attribute '{name}'")

    return node.attrib[name]

def get_attribute_optional(node: etree.Element, name: str) -> Optional[str]:
    return node.attrib.get(name, None)

def get_attribute_optional_int(node: etree.Element, name: str) -> Optional[int]:
    if not name in node.attrib: return None
    val = node.attrib[name]
    try:
        return int(val)
    except ValueError as ex:
        raise type(ex)(f"Could not convert attribute '{name}' of node '{node.tag}' in line {node.sourceline} to integer") from ex

def get_attribute_int(node: etree.Element, name: str) -> int:
    val = get_attribute_optional_int(node, name)
    if val is None:
        raise KeyError(f"Node '{node.tag}' in line {node.sourceline} is missing attribute '{name}'")
    return val

def get_attribute_optional_float(node: etree.Element, name: str) -> Optional[float]:
    if not name in node.attrib: return None
    val = node.attrib[name]
    try:
        return float(val)
    except ValueError as ex:
        raise type(ex)(f"Could not convert attribute '{name}' of node '{node.tag}' in line {node.sourceline} to float") from ex

def get_attribute_float(node: etree.Element, name: str) -> float:
    val = get_attribute_optional_float(node, name)
    if val is None:
        raise KeyError(f"Node '{node.tag}' in line {node.sourceline} is missing attribute '{name}'")
    return val


# Homebank Data classes

@dataclass
class Payee:
    name: str

    def __repr__(self) -> str:
        return f"Payee(\"{self.name}\")"

    @staticmethod
    def _from_node(node: etree.Element) -> Tuple[int, "Payee"]:
        key = get_attribute_int(node, "key")
        name = get_attribute(node, "name")
        return key, Payee(name)


@dataclass
class Account:
    name: str

    def __repr__(self) -> str:
        return f"Account(\"{self.name}\")"

    @staticmethod
    def _from_node(node: etree.Element) -> Tuple[int, "Account"]:
        key = get_attribute_int(node, "key")
        name = get_attribute(node, "name")
        return key, Account(name)


@dataclass
class Category:
    name: str
    parent: Optional[Category] = None

    def __repr__(self) -> str:
        if self.parent is not None:
            return f"Category(\"{self.name}\", {repr(self.parent)})"
        else:
            return f"Category(\"{self.name}\", None)"

    @staticmethod
    def _from_node(node: etree.Element, categories: Dict[int, "Category"]) -> Tuple[int, "Category"]:
        key: int = get_attribute_int(node, "key")
        name: str = get_attribute(node, "name")
        parent: Optional[int] = get_attribute_optional_int(node, "parent")
        parent_class: Optional["Category"] = None
        if parent is not None:
            if parent not in categories:
                raise KeyError(f"Parent category {parent} for category '{name}' not found")
            parent_class = categories[parent]
        return key, Category(name, parent_class)


@dataclass
class Operation:
    date: datetime.date
    amount: float
    account: Account
    dst_account: Optional[Account]
    payee: Optional[Payee]
    category: Optional[Category]
    wording: Optional[str]

    @staticmethod
    def _from_node(node: etree.Element, accounts: Dict[int, Account], payees: Dict[int, Payee], categories: Dict[int, Category]) -> "Operation":
        date_i = get_attribute_int(node, "date")
        date = datetime.date(1,1,1) + datetime.timedelta(days=date_i-1)
        amount = get_attribute_float(node, "amount")
        def find_account(account_i: int):
            if not account_i in accounts:
                raise KeyError(f"Account {account} for operation ({date}, {amount}) not found")
            return accounts[account_i]
        account_i = get_attribute_int(node, "account")
        account = find_account(account_i)
        dst_account_i = get_attribute_optional_int(node, "dst_account")
        dst_account: Optional[Account] = None
        if dst_account_i is not None:
            dst_account = find_account(dst_account_i)
        payee_i = get_attribute_optional_int(node, "payee")
        payee: Optional[Payee] = None
        if payee_i is not None:
            payee = payees[payee_i]
        category_i = get_attribute_optional_int(node, "category")
        category: Optional[Category] = None
        if category_i is not None:
            category = categories[category_i]
        wording = get_attribute_optional(node, "wording")
        return Operation(date, amount, account, dst_account, payee, category, wording)


class Homebank:
    def __init__(self, path: str):
        self.payees: Dict[int, Payee] = {}
        self.accounts: Dict[int, Account] = {}
        self.categories: Dict[int, Category] = {}
        self.operations: List[Operation] = []

        root = etree.parse(path).getroot()

        for node in root:
            if node.tag == "account":
                key, account = Account._from_node(node)
                self.accounts[key] = account
            elif node.tag == "pay":
                key, payee = Payee._from_node(node)
                self.payees[key] = payee
            elif node.tag == "cat":
                key, category = Category._from_node(node, self.categories)
                self.categories[key] = category
            elif node.tag == "ope":
                operation = Operation._from_node(node, self.accounts, self.payees, self.categories)
                self.operations.append(operation)