feat: Add a way to allow to get budgets information (#61)

Allows to get the budget information based either on the category or month. Allows to also create those budgets using the Python API.

---------

Co-authored-by: Brunno Vanelli <brunnovanelli@gmail.com>

actual/database.py

@@ -130,7 +130,8 @@ def before_flush(sess, flush_context, instances):
     @event.listens_for(session, "after_commit")
     @event.listens_for(session, "after_soft_rollback")
     def after_commit_or_rollback(
-        sess, previous_transaction=None  # noqa: previous_transaction needed for soft rollback
+        sess,
+        previous_transaction=None,  # noqa: previous_transaction needed for soft rollback
     ):
         if sess.info.get("messages"):
             del sess.info["messages"]
@@ -272,6 +273,12 @@ class Categories(BaseModel, table=True):
     tombstone: Optional[int] = Field(default=None, sa_column=Column("tombstone", Integer, server_default=text("0")))
     goal_def: Optional[str] = Field(default=None, sa_column=Column("goal_def", Text, server_default=text("null")))
 
+    zero_budgets: "ZeroBudgets" = Relationship(
+        back_populates="category",
+        sa_relationship_kwargs={
+            "primaryjoin": "and_(ZeroBudgets.category_id == Categories.id)",
+        },
+    )
     transactions: List["Transactions"] = Relationship(
         back_populates="category",
         sa_relationship_kwargs={
@@ -618,17 +625,37 @@ class ZeroBudgetMonths(SQLModel, table=True):
     buffered: Optional[int] = Field(default=None, sa_column=Column("buffered", Integer, server_default=text("0")))
 
 
-class ZeroBudgets(SQLModel, table=True):
+class ZeroBudgets(BaseModel, table=True):
     __tablename__ = "zero_budgets"
 
     id: Optional[str] = Field(default=None, sa_column=Column("id", Text, primary_key=True))
     month: Optional[int] = Field(default=None, sa_column=Column("month", Integer))
-    category: Optional[str] = Field(default=None, sa_column=Column("category", Text))
+    category_id: Optional[str] = Field(default=None, sa_column=Column("category", ForeignKey("categories.id")))
     amount: Optional[int] = Field(default=None, sa_column=Column("amount", Integer, server_default=text("0")))
     carryover: Optional[int] = Field(default=None, sa_column=Column("carryover", Integer, server_default=text("0")))
     goal: Optional[int] = Field(default=None, sa_column=Column("goal", Integer, server_default=text("null")))
     long_goal: Optional[int] = Field(default=None, sa_column=Column("long_goal", Integer, server_default=text("null")))
 
+    category: "Categories" = Relationship(
+        back_populates="zero_budgets",
+        sa_relationship_kwargs={
+            "uselist": False,
+            "primaryjoin": "and_(ZeroBudgets.category_id == Categories.id, Categories.tombstone == 0)",
+        },
+    )
+
+    def get_date(self) -> datetime.date:
+        return datetime.datetime.strptime(str(self.month), "%Y%m").date()
+
+    def set_date(self, date: datetime.date):
+        self.month = int(datetime.date.strftime(date, "%Y%m"))
+
+    def set_amount(self, amount: Union[decimal.Decimal, int, float]):
+        self.amount = int(round(amount * 100))
+
+    def get_amount(self) -> decimal.Decimal:
+        return decimal.Decimal(self.amount) / decimal.Decimal(100)
+
 
 class PendingTransactions(SQLModel, table=True):
     __tablename__ = "pending_transactions"

actual/queries.py

@@ -23,6 +23,7 @@
     Rules,
     Schedules,
     Transactions,
+    ZeroBudgets,
 )
 from actual.exceptions import ActualError
 from actual.rules import Action, Condition, Rule, RuleSet
@@ -109,9 +110,9 @@ def match_transaction(
     imported_id: str | None = None,
     already_matched: typing.List[Transactions] = None,
 ) -> typing.Optional[Transactions]:
-    """Matches a transaction with another transaction based on the fuzzy matching described at `reconcileTransactions`:
-
-    https://github.com/actualbudget/actual/blob/b192ad955ed222d9aa388fe36557b39868029db4/packages/loot-core/src/server/accounts/sync.ts#L347
+    """Matches a transaction with another transaction based on the fuzzy matching described at
+    [`reconcileTransactions`](
+    https://github.com/actualbudget/actual/blob/b192ad955ed222d9aa388fe36557b39868029db4/packages/loot-core/src/server/accounts/sync.ts#L347)
 
     The matches, from strongest to the weakest are defined as follows:
 
@@ -550,6 +551,71 @@ def get_or_create_account(s: Session, name: str | Accounts) -> Accounts:
     return account
 
 
+def get_budgets(
+    s: Session, month: datetime.date = None, category: str | Categories = None
+) -> typing.Sequence[ZeroBudgets]:
+    """
+    Returns a list of all available budgets.
+
+    :param s: session from Actual local database.
+    :param month: month to get budgets for, as a date for that month. Use `datetime.date.today()` if you want the budget
+                  for current month
+    :param category: category to filter for the budget. By default, the query looks for all budgets.
+    :return: list of budgets
+    """
+    query = select(ZeroBudgets).options(joinedload(ZeroBudgets.category))
+    if month:
+        month_filter = int(datetime.date.strftime(month, "%Y%m"))
+        query = query.filter(ZeroBudgets.month == month_filter)
+    if category:
+        category = get_category(s, category)
+        if not category:
+            raise ActualError("Category is provided but does not exist.")
+        query = query.filter(ZeroBudgets.category_id == category.id)
+    return s.exec(query).unique().all()
+
+
+def get_budget(s: Session, month: datetime.date, category: str | Categories) -> typing.Optional[ZeroBudgets]:
+    """
+    Gets an existing budget by category name, returns `None` if not found.
+
+    :param s: session from Actual local database.
+    :param month: month to get budgets for, as a date for that month. Use `datetime.date.today()` if you want the budget
+                  for current month.
+    :param category: category to filter for the budget.
+    :return: return budget matching the month and category. If not found, returns `None`.
+    """
+    budgets = get_budgets(s, month, category)
+    return budgets[0] if budgets else None
+
+
+def create_budget(
+    s: Session, month: datetime.date, category: str | Categories, amount: decimal.Decimal | float | int = 0.0
+) -> ZeroBudgets:
+    """
+    Gets an existing budget based on the month and category. If it already exists, the amount will be replaced by
+    the new amount.
+
+    :param s: session from Actual local database.
+    :param month: month to get budgets for, as a date for that month. Use `datetime.date.today()` if you want the budget
+                  for current month.
+    :param category: category to filter for the budget.
+    :param amount: amount for the budget.
+    :return: return budget matching the month and category, and assigns the amount to the budget. If not found, creates
+             a new budget.
+    """
+    budget = get_budget(s, month, category)
+    if budget:
+        budget.set_amount(amount)
+        return budget
+    category = get_category(s, category)
+    budget = ZeroBudgets(id=str(uuid.uuid4()), category_id=category.id)
+    budget.set_date(month)
+    budget.set_amount(amount)
+    s.add(budget)
+    return budget
+
+
 def create_transfer(
     s: Session,
     date: datetime.date,

tests/test_database.py

@@ -8,11 +8,13 @@
 from actual.database import Notes
 from actual.queries import (
     create_account,
+    create_budget,
     create_rule,
     create_splits,
     create_transaction,
     create_transfer,
     get_accounts,
+    get_budgets,
     get_or_create_category,
     get_or_create_payee,
     get_ruleset,
@@ -181,6 +183,26 @@ def test_rule_insertion_method(session):
     assert str(rs) == "If all of these conditions match 'date' isapprox '2024-01-02' then set 'cleared' to 'True'"
 
 
+def test_budgets(session):
+    # insert a budget
+    category = get_or_create_category(session, "Expenses")
+    session.commit()
+    create_budget(session, date(2024, 10, 7), category, 10.0)
+    assert len(get_budgets(session)) == 1
+    assert len(get_budgets(session, date(2024, 10, 1))) == 1
+    assert len(get_budgets(session, date(2024, 10, 1), category)) == 1
+    assert len(get_budgets(session, date(2024, 9, 1))) == 0
+    budget = get_budgets(session)[0]
+    assert budget.get_amount() == 10.0
+    assert budget.get_date() == date(2024, 10, 1)
+    # get a budget that already exists, but re-set it
+    create_budget(session, date(2024, 10, 7), category, 20.0)
+    assert budget.get_amount() == 20.0
+    # test if it fails if category does not exist
+    with pytest.raises(ActualError, match="Category is provided but does not exist"):
+        get_budgets(session, category="foo")
+
+
 def test_normalize_payee():
     assert normalize_payee("   mY paYeE ") == "My Payee"
     assert normalize_payee("  ", raw_payee_name=True) == ""