Files
WYGIWYH/app/apps/insights/utils/sankey.py
2025-02-15 00:48:59 -03:00

103 lines
3.5 KiB
Python

from django.db.models import Sum
from django.utils.translation import gettext_lazy as _
from decimal import Decimal
from typing import Dict, List, TypedDict
class SankeyNode(TypedDict):
name: str
class SankeyFlow(TypedDict):
from_node: str
to_node: str
flow: float
currency: Dict
original_amount: float
percentage: float
def generate_sankey_data(transactions_queryset):
"""
Generates Sankey diagram data from transaction queryset.
Uses a 1-5 scale for flows based on percentages.
"""
nodes: Dict[str, SankeyNode] = {}
flows: List[SankeyFlow] = []
# Aggregate transactions
income_data = {} # {(category, currency, account) -> amount}
expense_data = {} # {(category, currency, account) -> amount}
total_amount = Decimal("0")
for transaction in transactions_queryset:
currency = transaction.account.currency
account = transaction.account
category = transaction.category or _("Uncategorized")
key = (category, currency, account)
if transaction.type == "IN":
income_data[key] = income_data.get(key, Decimal("0")) + transaction.amount
else:
expense_data[key] = expense_data.get(key, Decimal("0")) + transaction.amount
total_amount += transaction.amount
# Function to add flow
def add_flow(from_node, to_node, amount, currency):
percentage = (amount / total_amount) * 100 if total_amount else 0
scaled_flow = 1 + min(percentage / 20, 4) # Scale 1-5, capping at 100%
flows.append(
{
"from_node": from_node,
"to_node": to_node,
"flow": float(scaled_flow),
"currency": {
"code": currency.code,
"prefix": currency.prefix,
"suffix": currency.suffix,
"decimal_places": currency.decimal_places,
},
"original_amount": float(amount),
"percentage": float(percentage),
}
)
nodes[from_node] = {"name": from_node}
nodes[to_node] = {"name": to_node}
# Process income
for (category, currency, account), amount in income_data.items():
category_name = f"{category} ({currency.code})"
account_name = f"{account.name} ({currency.code})"
add_flow(category_name, account_name, amount, currency)
# Process expenses
for (category, currency, account), amount in expense_data.items():
category_name = f"{category} ({currency.code})"
account_name = f"{account.name} ({currency.code})"
add_flow(account_name, category_name, amount, currency)
# Calculate and add savings flows
savings_data = {} # {(account, currency) -> amount}
for (category, currency, account), amount in income_data.items():
key = (account, currency)
savings_data[key] = savings_data.get(key, Decimal("0")) + amount
for (category, currency, account), amount in expense_data.items():
key = (account, currency)
savings_data[key] = savings_data.get(key, Decimal("0")) - amount
for (account, currency), amount in savings_data.items():
if amount > 0:
account_name = f"{account.name} ({currency.code})"
savings_name = f"{_('Savings')} ({currency.code})"
add_flow(account_name, savings_name, amount, currency)
return {
"nodes": list(nodes.values()),
"flows": flows,
"total_amount": float(total_amount),
}