diff --git a/pfbudget/cli/runnable.py b/pfbudget/cli/runnable.py index 3891241..766baea 100644 --- a/pfbudget/cli/runnable.py +++ b/pfbudget/cli/runnable.py @@ -4,6 +4,7 @@ import re from pfbudget.common.types import Operation from pfbudget.core.categories import categorize_data +from pfbudget.db.model import Period from pfbudget.input.json import JsonParser from pfbudget.input.nordigen import NordigenInput from pfbudget.db.sqlite import DatabaseClient @@ -333,6 +334,12 @@ def category(parser: argparse.ArgumentParser, universal: argparse.ArgumentParser update.add_argument("category", nargs="+", type=str) update.add_argument("--group", nargs="?", type=str) + schedule = commands.add_parser("schedule", parents=[universal]) + schedule.set_defaults(op=Operation.CategorySchedule) + schedule.add_argument("category", nargs="+", type=str) + schedule.add_argument("period", nargs=1, choices=[e.value for e in Period]) + schedule.add_argument("--frequency", nargs=1, default=[1], type=int) + group = commands.add_parser("group", parents=[universal]) category_group(group, universal) diff --git a/pfbudget/common/types.py b/pfbudget/common/types.py index 5439d7b..2c30b9c 100644 --- a/pfbudget/common/types.py +++ b/pfbudget/common/types.py @@ -16,6 +16,7 @@ class Operation(Enum): CategoryAdd = auto() CategoryUpdate = auto() CategoryRemove = auto() + CategorySchedule = auto() GroupAdd = auto() GroupRemove = auto() diff --git a/pfbudget/core/manager.py b/pfbudget/core/manager.py index 4f9ea11..9dd2911 100644 --- a/pfbudget/core/manager.py +++ b/pfbudget/core/manager.py @@ -2,7 +2,7 @@ from pfbudget.input.input import Input from pfbudget.input.nordigen import NordigenClient from pfbudget.input.parsers import parse_data from pfbudget.db.client import DbClient -from pfbudget.db.model import Category, CategoryGroup +from pfbudget.db.model import Category, CategoryGroup, CategorySchedule from pfbudget.common.types import Operation from pfbudget.core.categorizer import Categorizer from pfbudget.utils import convert @@ -67,6 +67,21 @@ class Manager: [Category(name=category) for category in self.args["category"]] ) + case Operation.CategorySchedule: + assert ( + "period" in self.args and "frequency" in self.args + ), "Schedule not well defined" + + with self.db.session() as session: + session.updateschedules( + [Category(name=category) for category in self.args["category"]], + CategorySchedule( + recurring=True, + period=self.args["period"][0], + period_multiplier=self.args["frequency"][0], + ), + ) + case Operation.GroupAdd: with self.db.session() as session: for group in self.args["group"]: diff --git a/pfbudget/db/client.py b/pfbudget/db/client.py index 88bc213..d2da511 100644 --- a/pfbudget/db/client.py +++ b/pfbudget/db/client.py @@ -1,7 +1,16 @@ +from copy import deepcopy +from dataclasses import asdict from sqlalchemy import create_engine, delete, select, update +from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session, joinedload, selectinload -from pfbudget.db.model import Bank, Category, CategoryGroup, Transaction +from pfbudget.db.model import ( + Bank, + Category, + CategoryGroup, + CategorySchedule, + Transaction, +) # import logging @@ -91,6 +100,30 @@ class DbClient: ) self.__session.execute(stmt) + def updateschedules( + self, categories: list[Category], schedule: CategorySchedule + ): + stmt = insert(CategorySchedule).values( + [ + dict( + name=cat.name, + recurring=schedule.recurring, + period=schedule.period, + period_multiplier=schedule.period_multiplier, + ) + for cat in categories + ] + ) + stmt = stmt.on_conflict_do_update( + index_elements=[CategorySchedule.name], + set_=dict( + recurring=stmt.excluded.recurring, + period=stmt.excluded.period, + period_multiplier=stmt.excluded.period_multiplier, + ), + ) + self.__session.execute(stmt) + def addcategorygroup(self, group: CategoryGroup): self.__session.add(group) diff --git a/pfbudget/db/model.py b/pfbudget/db/model.py index d5704e5..17470b7 100644 --- a/pfbudget/db/model.py +++ b/pfbudget/db/model.py @@ -107,12 +107,10 @@ class Category(Base): rules: Mapped[Optional[set[CategoryRule]]] = relationship( cascade="all, delete-orphan", passive_deletes=True ) - schedule: Mapped[CategorySchedule] = relationship() + schedule: Mapped[CategorySchedule] = relationship(back_populates="category") def __repr__(self) -> str: - return ( - f"Category(name={self.name}, group={self.group}, #rules={len(self.rules)})" - ) + return f"Category(name={self.name}, group={self.group}, #rules={len(self.rules)}, schedule={self.schedule})" catfk = Annotated[ @@ -220,3 +218,12 @@ class CategorySchedule(Base): recurring: Mapped[bool] period: Mapped[Optional[scheduleperiod]] period_multiplier: Mapped[Optional[int]] + + category: Mapped[Category] = relationship(back_populates="schedule") + + def __repr__(self) -> str: + return ( + f"{self.name} schedule=Schedule(period={self.period}, multiplier={self.period_multiplier})" + if self.recurring + else f"{self.name} has no Schedule" + )