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 "",
         }