Finishes category schedule implementation

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)
 

pfbudget/common/types.py

@@ -16,6 +16,7 @@ class Operation(Enum):
     CategoryAdd = auto()
     CategoryUpdate = auto()
     CategoryRemove = auto()
+    CategorySchedule = auto()
     GroupAdd = auto()
     GroupRemove = auto()
 

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"]:

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)
 

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"
+        )