utility functions for category

src/domain_utils.py

@@ -0,0 +1,87 @@
+from datetime import date, timedelta
+from decimal import Decimal
+
+from sqlalchemy.orm import Session
+
+from src.models.category import Category, CategoryBudget, CategoryBudgetDb, CategoryDb
+
+
+def find_category_db(session: Session, category: Category | None) -> CategoryDb | None:
+    if category is None:
+        return None
+
+    def find_existing(node: Category) -> CategoryDb | None:
+        if node.parent is None:
+            parent_id = None
+        else:
+            parent_db = find_existing(node.parent)
+            if parent_db is None:
+                return None
+            parent_id = parent_db.id
+
+        return (
+            session.query(CategoryDb)
+            .filter_by(name=node.name, parent_id=parent_id)
+            .first()
+        )
+
+    return find_existing(category)
+
+
+def set_category_budget(session: Session, category_budget: CategoryBudget):
+    category_db: CategoryDb | None = find_category_db(session, category_budget.category)
+    if category_db is None:
+        raise ValueError(
+            f"No matching category in database for category: {category_budget.category.to_str()}"
+        )
+
+    new_start = category_budget.valid_from or date.min
+    new_end = category_budget.valid_to or date.max
+
+    existing_budgets = (
+        session.query(CategoryBudgetDb).filter_by(category_id=category_db.id).all()
+    )
+
+    def _add_row(amount_eur: Decimal, valid_from: date | None, valid_to: date | None):
+        session.add(
+            CategoryBudgetDb(
+                category_id=category_db.id,
+                amount_eur=amount_eur,
+                valid_from=valid_from,
+                valid_to=valid_to,
+            )
+        )
+
+    for existing_budget in existing_budgets:
+        existing_start = existing_budget.valid_from or date.min
+        existing_end = existing_budget.valid_to or date.max
+
+        overlap_start = max(new_start, existing_start)
+        overlap_end = min(new_end, existing_end)
+
+        if overlap_start > overlap_end:
+            continue  # no overlap
+
+        session.delete(existing_budget)
+
+        if existing_start < new_start:
+            # Keep existing_start but cut existing end before new range
+            _add_row(
+                amount_eur=existing_budget.amount_eur,
+                valid_from=existing_budget.valid_from,
+                valid_to=overlap_start - timedelta(days=1),
+            )
+        if existing_end > new_end:
+            # Cut exsisting start to right after the new range and keep existing_end
+            _add_row(
+                amount_eur=existing_budget.amount_eur,
+                valid_from=overlap_end + timedelta(days=1),
+                valid_to=existing_budget.valid_to,
+            )
+
+    # Add the new budget
+    _add_row(
+        amount_eur=category_budget.amount_eur,
+        valid_from=category_budget.valid_from,
+        valid_to=category_budget.valid_to,
+    )

src/models/__init__.py

@@ -33,7 +33,7 @@ def __dir__() -> list[str]:
 if TYPE_CHECKING:
     from src.models.b4w import B4wExport as B4wExport
     from src.models.category import Category as Category
-    from src.models.category_budget import CategoryBudgetDb as CategoryBudgetDb
+    from src.models.category import CategoryBudgetDb as CategoryBudgetDb
     from src.models.comdirect import ComdirectCreditCard as ComdirectCreditCard
     from src.models.comdirect import ComdirectDepot as ComdirectDepot
     from src.models.comdirect import ComdirectGiro as ComdirectGiro

src/models/category.py

@@ -1,9 +1,12 @@
-from pydantic import Field, computed_field
-from sqlalchemy import ForeignKey, String, UniqueConstraint
+from datetime import date
+from decimal import Decimal
+from typing import Self
+
+from pydantic import BaseModel, Field, computed_field, model_validator
+from sqlalchemy import Date, ForeignKey, Numeric, String, UniqueConstraint
 from sqlalchemy.orm import Mapped, Session, mapped_column, relationship
 
 from src.db_utils import upsert
-from src.models.category_budget import CategoryBudgetDb
 from src.models.common import BaseModelDb, ExportBaseModel, hashify
 
 
@@ -74,3 +77,44 @@ def upsert_category(
         do_update=False,
         do_commit=do_commit,
     )
+
+
+class CategoryBudget(BaseModel):
+    category: Category
+    amount_eur: Decimal = Field(gt=0, decimal_places=2)
+    valid_from: date | None = Field(
+        description="Start of validity. None means unbounded"
+    )
+    valid_to: date | None = Field(description="End of validity. None means unbounded")
+
+    @model_validator(mode="after")
+    def validate_date_range(self) -> Self:
+        if (
+            self.valid_from is not None
+            and self.valid_to is not None
+            and self.valid_from > self.valid_to
+        ):
+            raise ValueError(
+                f"`valid_from` ({self.valid_from}) and valid_to ({self.valid_to}) are nor in order"
+            )
+        return self
+
+
+class CategoryBudgetDb(BaseModelDb):
+    """Budget amount for a category, optionally bounded to a time window.
+
+    Satellite table to ``category``: one category may have several budgets
+    (e.g. one per month). ``valid_from``/``valid_to`` are nullable; ``None``
+    means unbounded on that side.
+    """
+
+    __tablename__ = "category_budget"
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    category_id: Mapped[int] = mapped_column(ForeignKey("category.id"), nullable=False)
+    amount_eur: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False)
+    valid_from: Mapped[date | None] = mapped_column(Date, nullable=True, default=None)
+    valid_to: Mapped[date | None] = mapped_column(Date, nullable=True, default=None)
+
+    # relationship
+    category = relationship("CategoryDb", back_populates="budgets")

src/models/category_budget.py

@@ -1,27 +0,0 @@
-from datetime import date
-from decimal import Decimal
-
-from sqlalchemy import Date, ForeignKey, Numeric
-from sqlalchemy.orm import Mapped, mapped_column, relationship
-
-from src.models.common import BaseModelDb
-
-
-class CategoryBudgetDb(BaseModelDb):
-    """Budget amount for a category, optionally bounded to a time window.
-
-    Satellite table to ``category``: one category may have several budgets
-    (e.g. one per month). ``valid_from``/``valid_to`` are nullable; ``None``
-    means unbounded on that side.
-    """
-
-    __tablename__ = "category_budget"
-
-    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
-    category_id: Mapped[int] = mapped_column(ForeignKey("category.id"), nullable=False)
-    amount_eur: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False)
-    valid_from: Mapped[date | None] = mapped_column(Date, nullable=True, default=None)
-    valid_to: Mapped[date | None] = mapped_column(Date, nullable=True, default=None)
-
-    # relationship
-    category = relationship("CategoryDb", back_populates="budgets")