some fixes
This commit is contained in:
+179
-38
@@ -7,6 +7,11 @@ from src.models.category import Category, CategoryBudget, CategoryBudgetDb, Cate
|
||||
|
||||
|
||||
def find_category_db(session: Session, category: Category | None) -> CategoryDb | None:
|
||||
"""Find a CategoryDb in the database given a Pydantic Category object.
|
||||
|
||||
Recursively finds parent categories to build the full hierarchy path.
|
||||
Returns None if any part of the hierarchy doesn't exist.
|
||||
"""
|
||||
if category is None:
|
||||
return None
|
||||
|
||||
@@ -29,59 +34,195 @@ def find_category_db(session: Session, category: Category | None) -> CategoryDb
|
||||
|
||||
|
||||
def set_category_budget(session: Session, category_budget: CategoryBudget):
|
||||
category_db: CategoryDb | None = find_category_db(session, category_budget.category)
|
||||
"""Set or update a category budget, handling overlaps correctly.
|
||||
|
||||
This function handles three cases:
|
||||
1. No overlap: Add the new budget alongside existing ones
|
||||
2. Partial overlap: Split the existing budget and add appropriate parts + new budget
|
||||
3. Full overlap: Replace the existing budget entirely with the new one
|
||||
|
||||
Args:
|
||||
session: SQLAlchemy session
|
||||
category_budget: The budget to set/update
|
||||
|
||||
Raises:
|
||||
ValueError: If the category doesn't exist in the database
|
||||
"""
|
||||
# Find the category in the database
|
||||
category_db = 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()
|
||||
)
|
||||
new_start = category_budget.valid_from
|
||||
new_end = category_budget.valid_to
|
||||
|
||||
# Helper to add a budget row
|
||||
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,
|
||||
)
|
||||
return 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
|
||||
# Helper to check if two date ranges overlap
|
||||
def _ranges_overlap(start1: date | None, end1: date | None,
|
||||
start2: date | None, end2: date | None) -> bool:
|
||||
"""Check if two date ranges overlap.
|
||||
|
||||
Two bounded ranges overlap if start2 < end1 and start1 < end2.
|
||||
An unbounded range (None, None) only overlaps with another fully unbounded range.
|
||||
Partially unbounded ranges (e.g., valid_from=None but valid_to=date)
|
||||
only overlap if the bounded ranges themselves would overlap.
|
||||
"""
|
||||
# Both completely unbounded - always overlap
|
||||
if start1 is None and end1 is None and start2 is None and end2 is None:
|
||||
return True
|
||||
|
||||
# At least one range has both bounds None (truly unbounded)
|
||||
# It only overlaps with another truly unbounded range, not with bounded ranges
|
||||
if start1 is None and end1 is None or start2 is None and end2 is None:
|
||||
return False
|
||||
|
||||
# Both ranges are bounded - standard overlap check
|
||||
return start2 < end1 and start1 < end2
|
||||
|
||||
overlap_start = max(new_start, existing_start)
|
||||
overlap_end = min(new_end, existing_end)
|
||||
# Helper to find overlapping sub-range
|
||||
def _overlap_range(start1: date | None, end1: date | None,
|
||||
start2: date | None, end2: date | None) -> tuple[date | None, date | None]:
|
||||
"""Return the overlapping sub-range of two ranges."""
|
||||
|
||||
# Handle completely unbounded cases
|
||||
if start1 is None and end1 is None:
|
||||
# First range is completely unbounded - overlap is second range
|
||||
return (start2, end2) if (start2 is not None or end2 is not None) else (None, None)
|
||||
|
||||
if start2 is None and end2 is None:
|
||||
# Second range is completely unbounded - overlap is first range
|
||||
return (start1, end1) if (start1 is not None or end1 is not None) else (None, None)
|
||||
|
||||
# Initialize with defaults
|
||||
o_start = None
|
||||
o_end = None
|
||||
|
||||
# Handle one-sided unbounded cases
|
||||
if start1 is not None and end1 is not None:
|
||||
# Range 1 is bounded, range 2 may be unbounded
|
||||
if start2 is None: # New starts at beginning, so overlap starts at existing_start
|
||||
o_start = start1
|
||||
else:
|
||||
o_start = max(start1, start2)
|
||||
|
||||
if end2 is None: # New ends at end of time, so overlap ends at existing_end
|
||||
o_end = end1
|
||||
elif o_start is not None:
|
||||
o_end = min(end1, end2)
|
||||
|
||||
if start2 is not None and end2 is not None:
|
||||
# Range 2 is bounded, range 1 may be unbounded
|
||||
if start1 is None: # Existing starts at beginning
|
||||
o_start = start2
|
||||
elif o_start is None:
|
||||
o_start = max(start1, start2)
|
||||
|
||||
if end1 is None: # Existing ends at end of time
|
||||
o_end = end2
|
||||
elif o_start is not None:
|
||||
o_end = min(end1, end2)
|
||||
|
||||
# Check for actual overlap
|
||||
if o_start is None or o_start > o_end:
|
||||
return (None, None) # No overlap
|
||||
|
||||
return (o_start, o_end)
|
||||
|
||||
if overlap_start > overlap_end:
|
||||
continue # no overlap
|
||||
existing_budgets = (
|
||||
session.query(CategoryBudgetDb)
|
||||
.filter_by(category_id=category_db.id)
|
||||
.all()
|
||||
)
|
||||
|
||||
session.delete(existing_budget)
|
||||
def _remove_and_split_budget(budget: CategoryBudgetDb):
|
||||
"""Remove the overlapping parts of an existing budget and add back non-overlapping parts."""
|
||||
# Check for overlap with new range
|
||||
if not _ranges_overlap(new_start, new_end, budget.valid_from, budget.valid_to):
|
||||
row = _add_row(
|
||||
amount_eur=budget.amount_eur,
|
||||
valid_from=budget.valid_from,
|
||||
valid_to=budget.valid_to,
|
||||
)
|
||||
session.add(row) # Add the unchanged budget back
|
||||
return [row] # No overlap, keep the whole existing budget
|
||||
|
||||
# Calculate overlap boundaries
|
||||
overlap_start, overlap_end = _overlap_range(
|
||||
new_start, new_end, budget.valid_from, budget.valid_to
|
||||
)
|
||||
|
||||
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,
|
||||
if overlap_start is None or overlap_end is None:
|
||||
# Full unbounded overlap - replace entirely with new budget
|
||||
row = _add_row(category_budget.amount_eur, category_budget.valid_from, category_budget.valid_to)
|
||||
session.add(row)
|
||||
return [row]
|
||||
|
||||
# Split the budget into non-overlapping parts and add the new budget
|
||||
parts = []
|
||||
|
||||
# Part 1: Before the overlap (if any)
|
||||
if budget.valid_from is not None and budget.valid_from < overlap_start:
|
||||
row = _add_row(
|
||||
amount_eur=budget.amount_eur,
|
||||
valid_from=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,
|
||||
parts.append(row)
|
||||
|
||||
# Part 2: After the overlap (if any)
|
||||
if budget.valid_to is not None and overlap_end < budget.valid_to:
|
||||
row = _add_row(
|
||||
amount_eur=budget.amount_eur,
|
||||
valid_from=overlap_end + timedelta(days=1),
|
||||
valid_to=existing_budget.valid_to,
|
||||
valid_to=budget.valid_to,
|
||||
)
|
||||
parts.append(row)
|
||||
|
||||
# Add the new budget
|
||||
row = _add_row(
|
||||
amount_eur=category_budget.amount_eur,
|
||||
valid_from=category_budget.valid_from,
|
||||
valid_to=category_budget.valid_to,
|
||||
)
|
||||
parts.append(row)
|
||||
|
||||
# Add the new budget
|
||||
_add_row(
|
||||
amount_eur=category_budget.amount_eur,
|
||||
valid_from=category_budget.valid_from,
|
||||
valid_to=category_budget.valid_to,
|
||||
)
|
||||
# Add all parts to session
|
||||
for row in parts:
|
||||
session.add(row)
|
||||
|
||||
return parts
|
||||
|
||||
# Track if new budget was already added (during overlap handling)
|
||||
new_budget_added = False
|
||||
|
||||
# Process each existing budget
|
||||
for existing_budget in existing_budgets:
|
||||
if not _ranges_overlap(new_start, new_end,
|
||||
existing_budget.valid_from, existing_budget.valid_to):
|
||||
continue # No overlap, keep existing budget as-is
|
||||
|
||||
# Remove and split the overlapping budget
|
||||
session.delete(existing_budget)
|
||||
|
||||
# Add back non-overlapping parts plus the new budget
|
||||
split_parts = _remove_and_split_budget(existing_budget)
|
||||
new_budget_added = True # Mark that we already added the new budget
|
||||
|
||||
# Always add the new budget unless it was already added during overlap handling
|
||||
if not new_budget_added:
|
||||
row = _add_row(
|
||||
amount_eur=category_budget.amount_eur,
|
||||
valid_from=category_budget.valid_from,
|
||||
valid_to=category_budget.valid_to,
|
||||
)
|
||||
session.add(row)
|
||||
|
||||
Reference in New Issue
Block a user