Added base schema for Marshmallow dataclasses.

This commit is contained in:
Fabio Manganiello 2023-03-26 03:47:44 +02:00
parent 7bbae55e44
commit 026662f6b6
Signed by: blacklight
GPG key ID: D90FBA7F76362774

View file

@ -0,0 +1,56 @@
from datetime import date, datetime
from uuid import UUID
from marshmallow import (
EXCLUDE,
Schema,
fields,
post_dump,
)
from .. import Date, DateTime
class DataClassSchema(Schema):
"""
Base schema class for data classes that support Marshmallow schemas.
"""
TYPE_MAPPING = {
date: Date,
datetime: DateTime,
UUID: fields.UUID,
}
# pylint: disable=too-few-public-methods
class Meta:
"""
Exclude unknown fields.
"""
unknown = EXCLUDE
def _get_field(self, key: str) -> fields.Field:
"""
Returns the matching field by either name or data_key.
"""
if key in self.fields:
return self.fields[key]
matching_fields = [f for f in self.fields.values() if key == f.data_key]
assert (
len(matching_fields) == 1
), f'Could not find field {key} in {self.__class__.__name__}'
return matching_fields[0]
@post_dump
def post_dump(self, data: dict, **__) -> dict:
# Use data_key parameters only for load
new_data = {}
for key, value in data.items():
field = self._get_field(key)
new_data[field.name if field.data_key is not None else key] = value
return new_data