diff --git a/src/domain_utils.py b/src/domain_utils.py new file mode 100644 index 0000000..e7a78a3 --- /dev/null +++ b/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, + ) diff --git a/src/models/__init__.py b/src/models/__init__.py index 6fe7e34..3ae1afe 100644 --- a/src/models/__init__.py +++ b/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 diff --git a/src/models/category.py b/src/models/category.py index 3dec8a8..cdb17e7 100644 --- a/src/models/category.py +++ b/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") diff --git a/src/models/category_budget.py b/src/models/category_budget.py deleted file mode 100644 index 14354fc..0000000 --- a/src/models/category_budget.py +++ /dev/null @@ -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") diff --git a/tests/test_category_budget.py b/tests/test_category_budget.py index b54deeb..ff738a1 100644 --- a/tests/test_category_budget.py +++ b/tests/test_category_budget.py @@ -4,8 +4,7 @@ from decimal import Decimal from sqlalchemy import inspect from sqlalchemy.orm import Session -from src.models.category import CategoryDb -from src.models.category_budget import CategoryBudgetDb +from src.models.category import CategoryBudgetDb, CategoryDb def test_category_budget_schema(db_session: Session):