model-service/model_service/schemas/aimodel_schema.py

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

from .svmdetails_schema import SVMDetailsSchema


class AIModelSchema(SQLAlchemyAutoSchema):

    _svm_details_schema = SVMDetailsSchema(many=False, exclude=["aimodel"])

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

    details = fields.Method("pluck_details", dump_only=True)

    type = EnumField(AIModelType)

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

    def pluck_details(self, ai_model) -> Optional[dict]:
        if ai_model.details:
            return self._svm_details_schema.dump(ai_model.details[0])
        else:
            return None

    class Meta:
        model = AIModel
        include_relationships = True
        load_instance = True
        include_fk = True