import numpy as np

from presupuestov1.schema import OrmBaseModel
from pydantic import field_serializer


class Treemap(OrmBaseModel):
    gestion: int
    entidad: int
    nivel: str | None
    parent: int | None
    devengado: float | None

    @field_serializer('devengado')
    def do_round(self, v):
        if v is not None:
            return np.round(v, 2)


class TreemapObjeto(Treemap):
    objeto: str
    desc_objeto: str | None


class TreemapFinFun(Treemap):
    finfun: str
    desc_finfun: str | None


class TreemapActeco(Treemap):
    acteco: str
    desc_actividad: str | None