finalized helper functions for category trees; started writing tests

src/ingest/transactions.py

@@ -5,11 +5,21 @@ from sqlalchemy.dialects.sqlite import insert as sqlite_insert
 from sqlalchemy.orm import Session
 
 from src.models.b4w import B4wDb
+from src.models.category import Category, CategoryDb
 from src.models.comdirect import ComdirectGiroDb
 from src.models.common import BaseModelDb
 from src.models.transactions import Transaction, TransactionDb
 
 
+def _category_from_db(category_db: CategoryDb | None) -> Category | None:
+    if category_db is None:
+        return None
+    return Category(
+        name=category_db.name,
+        parent=_category_from_db(category_db.parent),
+    )
+
+
 def map_b4w_to_transaction(row: B4wDb) -> Transaction:
     """Project one staging row of ``import_b4w`` onto a canonical Transaction."""
     return Transaction(
@@ -24,7 +34,7 @@ def map_b4w_to_transaction(row: B4wDb) -> Transaction:
         text=row.text,
         other_name=row.other_name,
         other_iban=row.other_iban,
-        category_id=row.category_id,
+        category=_category_from_db(row.category),
     )
 
 
@@ -44,7 +54,11 @@ def map_comdirect_giro_to_transaction(row: ComdirectGiroDb) -> Transaction:
 
 
 def upsert_transaction(
-    session: Session, transaction: Transaction, do_commit: bool = True
+    session: Session,
+    transaction: Transaction,
+    *,
+    category_db: CategoryDb | None = None,
+    do_commit: bool = True,
 ) -> None:
     """Insert one canonical transaction, ignoring duplicates on ``hash_id``.
 
@@ -65,7 +79,7 @@ def upsert_transaction(
         "text": transaction.text,
         "other_name": transaction.other_name,
         "other_iban": transaction.other_iban,
-        "category_id": transaction.category_id,
+        "category_id": category_db.id if category_db is not None else None,
     }
     stmt = sqlite_insert(TransactionDb).values(parameters)
     _ = session.execute(stmt.prefix_with("OR IGNORE"))
@@ -92,6 +106,11 @@ def consolidate_transactions[StagingDbT: BaseModelDb](
         transaction = mapper(staging_row)
         if transaction.hash_id in existing_hashes:
             continue
-        upsert_transaction(session, transaction, do_commit=False)
+        upsert_transaction(
+            session,
+            transaction,
+            category_db=getattr(staging_row, "category", None),
+            do_commit=False,
+        )
     if do_commit:
         session.commit()

src/models/category.py

@@ -18,8 +18,7 @@ from sqlalchemy import (
 from sqlalchemy.orm import Mapped, Session, mapped_column, relationship
 
 from src.db_utils import upsert
-from src.models.common import Amount, AmountPos, BaseModelDb, ExportBaseModel, hashify
-from src.models.transactions import Transaction
+from src.models.common import AmountPos, BaseModelDb, ExportBaseModel, hashify
 
 
 class UTCDateTime(TypeDecorator[datetime]):
@@ -70,19 +69,15 @@ class Category(ExportBaseModel):
         else:
             return f"{self.parent.to_str()}:{self.name}"
 
-    def __str__(self):
-        return self.to_str()
+    def __str__(self) -> str:
+        return self.to_str() or ""
 
     def __repr__(self):
         return f'Category.from_str("{self.__str__()}")'
 
     def get_budget(self, date_valid: date | None) -> float:
         # TODO
-        ...
-
-    def get_tree_to_root(self) -> CategoryTree:
-        # TODO (maybe)
-        ...
+        return 0.0
 
 
 class CategoryDb(BaseModelDb):
@@ -237,67 +232,3 @@ class CategoryBudgetDb(BaseModelDb):
 
     # relationship
     category = relationship("CategoryDb", back_populates="budgets")
-
-
-class CategoryTree(BaseModel):
-    category: Category
-    total_amount: float | None = None
-    total_budget: float | None = None
-    sub_categories: list[CategoryTree] | None = None
-    transactions: list[Transaction] | None = None
-
-    def get_total_amount(self) -> float:
-        result = 0.0
-        if self.transactions:
-            sum(transaction.amount for transaction in self.transactions)
-        if self.sub_categories:
-            sum(sub_category.get_total_amount() for sub_category in self.sub_categories)
-        return result
-
-    def get_total_budget(self, date_valid: date) -> float:
-        result = 0.0
-        result += self.category.get_budget(date_valid)
-        if self.sub_categories:
-            result += sum(
-                category.category.get_budget(date_valid)
-                for category in self.sub_categories
-            )
-        return result
-
-    def __add__(self, other: CategoryTree):
-        if self.category != other.category:
-            return
-        sub_categories = []
-        for category1 in self.sub_categories or []:
-            for category2 in other.sub_categories or []:
-                if category1.category == category2.category:
-                    sub_categories.append(category1 + category2)
-                else:
-                    sub_categories.append(category1)
-                    sub_categories.append(category2)
-
-        return CategoryTree(
-            category=self.category,
-            total_amount=(self.total_amount or 0.0) + (other.total_amount or 0.0),
-            total_budget=(self.total_budget or 0.0) + (other.total_budget or 0.0),
-            sub_categories=sub_categories,
-            transactions=self.transactions + other.transactions,
-        )
-
-
-def build_category_tree(transactions: list[Transaction]) -> list[CategoryTree]:
-    # TODO: loop through flat list a build a category tree
-    result = []
-    for transaction in transactions:
-        bla = _transaction_to_tree(transaction)
-
-
-def _transaction_to_tree(transaction: Transaction) -> CategoryTree | None:
-    category = transaction.category
-    if not category:
-        return
-    result = CategoryTree(category=category, transactions=[transaction])
-    while category.parent is not None:
-        category = category.parent
-        result = CategoryTree(category=category, sub_categories=[result])
-    return result

src/models/category_tree.py

@@ -0,0 +1,96 @@
+from datetime import date
+
+from pydantic import BaseModel
+
+from src.models.category import Category
+from src.models.transactions import Transaction
+
+
+class CategoryTree(BaseModel):
+    category: Category | None
+    total_amount: float | None = None
+    total_budget: float | None = None
+    sub_categories: list[CategoryTree] | None = None
+    transactions: list[Transaction] | None = None
+
+    def get_total_amount(self) -> float:
+        result = 0.0
+        if self.transactions:
+            sum(transaction.amount for transaction in self.transactions)
+        if self.sub_categories:
+            sum(sub_category.get_total_amount() for sub_category in self.sub_categories)
+        return result
+
+    def get_total_budget(self, date_valid: date) -> float:
+        result = 0.0
+        if self.category is None:
+            return result
+        result += self.category.get_budget(date_valid)
+        if self.sub_categories:
+            result += sum(
+                (
+                    category.category.get_budget(date_valid)
+                    if category.category is not None
+                    else 0.0
+                )
+                for category in self.sub_categories
+            )
+        return result
+
+    def __add__(self, other: CategoryTree):
+        if self.category != other.category:
+            return
+
+        total_amount = (
+            (self.total_amount or 0.0) + (other.total_amount or 0.0)
+        ) or None
+        total_budget = (
+            (self.total_budget or 0.0) + (other.total_budget or 0.0)
+        ) or None
+
+        sub_categories = []
+        for category1 in self.sub_categories or []:
+            for category2 in other.sub_categories or []:
+                if category1.category == category2.category:
+                    sub_categories.append(category1 + category2)
+                else:
+                    sub_categories.append(category1)
+                    sub_categories.append(category2)
+
+        transactions = ((self.transactions or []) + (other.transactions or [])) or None
+
+        return CategoryTree(
+            category=self.category,
+            total_amount=total_amount,
+            total_budget=total_budget,
+            sub_categories=sub_categories,
+            transactions=transactions,
+        )
+
+
+def build_category_tree(transactions: list[Transaction]) -> list[CategoryTree]:
+    category_branches = map(_transaction_to_category_branch, transactions)
+    results: list = []
+    for category_branch in category_branches:
+        if not results:
+            results += [category_branch]
+            continue
+        if category_branch.category not in [result.category for result in results]:
+            results += [category_branch]
+        else:
+            for result in results:
+                if (category_branch.category or "") == (result.category or ""):
+                    result += category_branch
+                    break
+    return results
+
+
+def _transaction_to_category_branch(transaction: Transaction) -> CategoryTree:
+    category = transaction.category
+    result = CategoryTree(category=category, transactions=[transaction])
+    if not category:
+        return result
+    while category.parent is not None:
+        category = category.parent
+        result = CategoryTree(category=category, sub_categories=[result])
+    return result

src/models/common.py

@@ -40,6 +40,7 @@ def date2str(value: date | None, default="--", date_format="%Y-%m-%d") -> str:
         return value.strftime(date_format)
     return default
 
+
 type AmountPos = Annotated[Decimal, Field(decimal_places=2, gt=0)]
 type Amount = Annotated[Decimal, Field(decimal_places=2)]
 
@@ -64,7 +65,6 @@ type Csv2Dict = list[Csv2DictRow]
 type CsV2DictMulti = dict[str, Csv2Dict]
 
 
-
 class ExportBaseModel(BaseModel):
 
     @classmethod

src/models/transactions.py

@@ -1,5 +1,4 @@
 from datetime import date
-from decimal import Decimal
 
 from pydantic import Field, computed_field
 from sqlalchemy import Date, ForeignKey, Index, Integer, String, UniqueConstraint
@@ -29,9 +28,7 @@ class Transaction(ExportBaseModel):
     text: str | None = None
     other_name: str | None = None
     other_iban: str | None = None
-    # TODO: The Pydantic model should hold a category object, NOT the category id which is owned by the DB
-    category: Category | None
-    category_id: int | None = None # remove!
+    category: Category | None = None
 
     @computed_field  # type: ignore[prop-decorator]
     @property
@@ -58,7 +55,7 @@ class Transaction(ExportBaseModel):
             "text": f"{self.text or ''}",
             "other_name": f"{self.other_name or ''}",
             "other_iban": f"{self.other_iban or ''}",
-            "category_id": f"{self.category_id}",
+            "category": (self.category.to_str() if self.category else "") or "",
         }
 
 

tests/test_category_edge_cases.py

@@ -1,4 +1,5 @@
 from datetime import date
+from decimal import Decimal
 
 import pytest
 
@@ -12,24 +13,30 @@ class TestCategoryComputedFields:
     def test_hash_id_computed_field_simple(self):
         """Test hash_id for single-level category."""
         c = Category.from_str("single")
+        assert c
         assert hasattr(c, "hash_id")
         assert isinstance(c.hash_id, str)
         # Hash should be deterministic for same inputs
         c2 = Category.from_str("single")
+        assert c2
         assert c.hash_id == c2.hash_id
 
     def test_hash_id_computed_field_nested(self):
         """Test hash_id for multi-level category hierarchy."""
         c1 = Category.from_str("level1:level2:level3")
+        assert c1
         assert hasattr(c1, "hash_id")
 
         c2 = Category.from_str("level1:level2:level3")
+        assert c2
         assert c1.hash_id == c2.hash_id
 
     def test_hash_id_different_hierarchy(self):
         """Test that different hierarchies have different hash_ids."""
         c1 = Category.from_str("A:B")
         c2 = Category.from_str("A:C")
+        assert c1
+        assert c2
         # Different last level -> different hash_id
         assert c1.hash_id != c2.hash_id
 
@@ -40,22 +47,26 @@ class TestCategoryComputedFields:
 
         # Create another with same name but no parent
         orphan = Category.from_str("Budget")
+        assert orphan
 
         assert cat1.hash_id != orphan.hash_id
 
     def test_to_str_single_level(self):
         """Test to_str() for single-level category."""
         c = Category.from_str("single")
+        assert c
         assert c.to_str() == "single"
 
     def test_to_str_nested(self):
         """Test to_str() for nested category."""
         c = Category.from_str("parent:child:grandchild")
+        assert c
         assert c.to_str() == "parent:child:grandchild"
 
     def test_str_uses_to_str(self):
         """Test that __str__ delegates to to_str()."""
         c = Category.from_str("A:B:C")
+        assert c
         assert str(c) == c.to_str() == "A:B:C"
 
     def test_repr_simple(self):
@@ -79,7 +90,7 @@ class TestCategoryBudgetValidation:
         cat = CategoryFactory.build()
         budget = CategoryBudget(
             category=cat,
-            amount_eur=100.00,
+            amount_eur=Decimal("100.00"),
         )
         assert budget.valid_from is None
         assert budget.valid_to is None
@@ -90,7 +101,7 @@ class TestCategoryBudgetValidation:
         cat = CategoryFactory.build()
         budget = CategoryBudget(
             category=cat,
-            amount_eur=100.00,
+            amount_eur=Decimal("100.00"),
             valid_from=date(2026, 10, 1),
             valid_to=date(2026, 10, 31),
         )
@@ -105,7 +116,7 @@ class TestCategoryBudgetValidation:
         with pytest.raises(ValueError) as exc_info:
             CategoryBudget(
                 category=cat,
-                amount_eur=100.00,
+                amount_eur=Decimal("100.00"),
                 valid_from=date(2026, 10, 31),
                 valid_to=date(2026, 10, 1),
             )
@@ -120,7 +131,7 @@ class TestCategoryBudgetValidation:
         cat = CategoryFactory.build()
         budget = CategoryBudget(
             category=cat,
-            amount_eur=100.00,
+            amount_eur=Decimal("100.00"),
             valid_from=date(2026, 10, 1),
             valid_to=date(2026, 10, 1),
         )
@@ -133,7 +144,7 @@ class TestCategoryBudgetValidation:
         cat = CategoryFactory.build()
         budget = CategoryBudget(
             category=cat,
-            amount_eur=150.75,
+            amount_eur=Decimal("150.75"),
             valid_from=date(2026, 9, 1),
             valid_to=date(2026, 9, 30),
         )
@@ -142,7 +153,7 @@ class TestCategoryBudgetValidation:
         try:
             CategoryBudget(
                 category=cat,
-                amount_eur=200.00,
+                amount_eur=Decimal("200.00"),
                 valid_from=date(2026, 12, 31),
                 valid_to=date(2026, 12, 30),
             )

tests/test_category_tree.py

@@ -0,0 +1,47 @@
+from src.models.category import Category
+from src.models.category_tree import CategoryTree, _transaction_to_category_branch
+from src.models.transactions import Transaction
+from tests.factories import TransactionFactory
+
+
+def test_transaction_to_category_branch_simple():
+    transaction = TransactionFactory.build(category=Category.from_str("dummy category"))
+    result = _transaction_to_category_branch(transaction)
+    assert isinstance(result, CategoryTree)
+    assert result.category
+    assert result.category.name == "dummy category"
+    assert result.transactions
+    assert isinstance(result.transactions, list)
+    assert len(result.transactions) == 1
+    assert isinstance(result.transactions[0], Transaction)
+    assert result.transactions[0] == transaction
+
+
+def test_transaction_to_category_branch_two_category_levels():
+    transaction = TransactionFactory.build(category=Category.from_str("father:son"))
+    result = _transaction_to_category_branch(transaction)
+    assert isinstance(result, CategoryTree)
+    assert result.category
+    assert result.category.name == "father"
+    assert result.transactions is None
+    assert result.sub_categories
+    assert isinstance(result.sub_categories, list)
+    assert len(result.sub_categories) == 1
+    assert isinstance(result.sub_categories[0], CategoryTree)
+    assert isinstance(result.sub_categories[0].category, Category)
+    assert result.sub_categories[0].category.name == "son"
+    assert result.sub_categories[0].transactions
+    assert isinstance(result.sub_categories[0].transactions, list)
+    assert len(result.sub_categories[0].transactions) == 1
+    assert isinstance(result.sub_categories[0].transactions[0], Transaction)
+    assert result.sub_categories[0].transactions[0] == transaction
+
+
+def test_transaction_to_category_branch_no_category():
+    transaction = TransactionFactory.build(category=None)
+    result = _transaction_to_category_branch(transaction)
+    assert result.category is None
+    assert result.transactions
+    assert isinstance(result.transactions, list)
+    assert len(result.transactions) == 1
+    assert result.transactions[0] == transaction

tests/test_transactions.py

@@ -101,7 +101,7 @@ def test_map_comdirect_giro_to_transaction():
     assert transaction.amount == Decimal("32.43")
     assert transaction.other_name is None
     assert transaction.other_iban is None
-    assert transaction.category_id is None
+    assert transaction.category is None
 
 
 def test_consolidate_b4w_signed_amount_and_category(db_session: Session):
@@ -125,6 +125,9 @@ def test_consolidate_b4w_signed_amount_and_category(db_session: Session):
     assert transaction.amount == Decimal("-50.00")
     assert transaction.category_id == staging.category_id
     assert transaction.text == "Supermarket"
+    mapped = map_b4w_to_transaction(staging)
+    assert mapped.category is not None
+    assert mapped.category.to_str() == export.category.to_str()
 
 
 def test_ingest_b4w_e2e_idempotent(db_session: Session):