model-service/model_service/schemas/aimodel_schema.py

32 lines
854 B
Python
Raw Normal View History

2020-04-14 02:01:44 +02:00
#!/usr/bin/env python3
2020-10-01 22:48:00 +02:00
from typing import Optional
2020-04-14 02:01:44 +02:00
from marshmallow import fields
from marshmallow_sqlalchemy import ModelSchema
from marshmallow_enum import EnumField
from model import AIModel, AIModelType
2020-04-14 02:01:44 +02:00
2020-10-01 22:48:00 +02:00
from .svmdetails_schema import SVMDetailsSchema
2020-04-14 02:01:44 +02:00
2020-04-14 13:48:11 +02:00
class AIModelSchema(ModelSchema):
2020-04-14 17:26:33 +02:00
2020-10-01 22:48:00 +02:00
_svm_details_schema = SVMDetailsSchema(many=False, exclude=["aimodel"])
2020-04-14 17:26:33 +02:00
default = fields.Method("boolize_default", dump_only=True)
2020-10-01 22:48:00 +02:00
details = fields.Method("pluck_details", dump_only=True)
type = EnumField(AIModelType)
2020-04-14 17:26:33 +02:00
def boolize_default(self, ai_model) -> bool:
return bool(ai_model.default)
2020-10-01 22:48:00 +02:00
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
2020-04-14 02:01:44 +02:00
class Meta:
2020-04-14 13:48:11 +02:00
model = AIModel