#!/usr/bin/env python3
from marshmallow import fields
from marshmallow_sqlalchemy import ModelSchema
from marshmallow_enum import EnumField
from model import AIModel, AIModelType


class AIModelSchema(ModelSchema):

    default = fields.Method("boolize_default", dump_only=True)

    type = EnumField(AIModelType)

    def boolize_default(self, ai_model) -> bool:
        return bool(ai_model.default)

    class Meta:
        model = AIModel