"""Handle Model Serialization.""" from typing import Dict, Callable, Any, List, Optional, Type from atheneum import errors from atheneum.db import db class BaseSerializer: # pylint: disable=too-few-public-methods """Base Model serializer.""" def __init__(self, model: Type[db.Model]) -> None: """Initialize the base serializer.""" self._fields: Dict[str, Callable[[db.Model], Any]] = {} self.model = model def serialize(self, options: Optional[List[str]]) -> Any: """Convert Model field and factories to dicts.""" field_factories = self._serializers() if not options: options = list(field_factories.keys()) ret = {} for key in options: if key not in field_factories: raise errors.ValidationError( 'Invalid key: %r. Valid keys: %r.' % ( key, list(sorted(field_factories.keys())))) factory = field_factories[key] val = factory() if val is not None: ret[key] = val return ret def _serializers(self) -> Dict[str, Callable[[], Any]]: """Field definitions.""" raise NotImplementedError() _model_serializers: Dict[str, Type[BaseSerializer]] = {} def register_serializer( model_name: str, model_serializer: Type[BaseSerializer]) -> None: """Add a model to the serializer mapping.""" _model_serializers[model_name] = model_serializer def serialize_model(model_obj: db.Model) -> Any: """Lookup a Model and hand off to the serializer.""" try: return _model_serializers[ type(model_obj).__name__](model_obj).serialize(None) except KeyError: raise NotImplementedError( '{} has no registered serializers'.format(model_obj.__name__))