from rest_framework import serializers

from .models import BudgetGroup, BudgetLine


class BudgetGroupSerializer(serializers.ModelSerializer):
    total = serializers.SerializerMethodField()

    def get_total(self, obj):
        tot = 0
        for i in obj.budgetline_set.all():
            tot += i.amount
        return tot

    class Meta:
        model = BudgetGroup
        fields = "__all__"


class BudgetLineSerializer(serializers.ModelSerializer):
    class Meta:
        model = BudgetLine
        fields = "__all__"