marshmallow-sqlalchemy

marshmallow-sqlalchemy copied to clipboard

I ran into a problem. When I use instance, I get a field type of enum.

print(UserSchema().load({}, instance=user).data.status)
# <StatusType.active: 1>

But if I use data, I get a field type of string.

print(UserSchema().load({'status': 'active'}, instance=user).data.status)
# active

I have a solution.

from sqlalchemy.types import Enum

class EnumField(fields.Field):

    def __init__(self, *args, **kwargs):
        self.column = kwargs.get('column')
        super(EnumField, self).__init__(*args, **kwargs)

    def _serialize(self, value, attr, obj):
        field = super(EnumField, self)._serialize(value, attr, obj)
        return field.name if field else field

    def deserialize(self, value, attr=None, data=None):
        field = super(EnumField, self).deserialize(value, attr, data)
        if isinstance(field, str) and self.column is not None:
            return self.column.type.python_type[field]
        return field


class ExtendModelConverter(ModelConverter):
    ModelConverter.SQLA_TYPE_MAPPING[Enum] = EnumField

    def _add_column_kwargs(self, kwargs, column):
        super(ExtendModelConverter, self)._add_column_kwargs(kwargs, column)
        if hasattr(column.type, 'enums'):
            kwargs['column'] = column

But I am confused by one place, this is the decorator marshmallow.pre_load. I still need to work with string at this place.

What do you think about it? This problem exists? What are the options for solving it?

The above code snippet works but it reaches into ModelConverter and changes SQLA_TYPE_MAPPING. Instead it would be better to do:

class EnumField(marshmallow.fields.Field):
    def __init__(self, *args, **kwargs):
        self.column = kwargs.get('column')
        super(EnumField, self).__init__(*args, **kwargs)

    def _serialize(self, value, attr, obj):
        field = super(EnumField, self)._serialize(value, attr, obj)
        return field.name if field else field

    def deserialize(self, value, attr=None, data=None):
        field = super(EnumField, self).deserialize(value, attr, data)
        if isinstance(field, str) and self.column is not None:
            return self.column.type.python_type[field]
        return field


class ExtendModelConverter(ModelConverter):
    SQLA_TYPE_MAPPING = {
        **ModelConverter.SQLA_TYPE_MAPPING,
        Enum: EnumField,
    }

    def _add_column_kwargs(self, kwargs, column):
        super()._add_column_kwargs(kwargs, column)
        if hasattr(column.type, 'enums'):
            kwargs['column'] = column

And in your schema:

class MySchema(ma.ModelSchema):
    class Meta:
        model = models.MyModel
        model_converter = ExtendModelConverter

The above code from @wjdp works incredible. Would any maintainer be able to comment whether this might be accepted into the core project? Seems like a no brainer to be able to support enums.

Only change I would suggest is to maybe use field.value instead of field.name and self.column.type.python_type(field) to serialize/deserialize, so that the custom integer/string value of the enum class is used. SQLA uses the "left hand side" name in the db, but there are a variety of reasons you might want to serialize to something other than the python/sql representation.

E.g. your enum may have:

value_a = "value-a"
value_b = "value-b"

And you would likely prefer the hyphenated version for serialization if you're doing something URL-related (because of inconsistent handling of underscores). This could of course be configurable somehow.

Completely missed this but enum is implemented already. Just needs a custom serializer.

For anyone stumbling upon this in 2022+:

class EnumField(fields.Field):
    """Marshmallow field for SQLA enum type"""
    def _serialize(self, value: Any, attr: str, obj: Any, **kwargs):
        """Serialize an enum type to a string"""
        try:
            return value.name
        except AttributeError:
            return value

class ExtendModelConverter(ModelConverter):
    """Set up type overrides for UUID and enum."""

    SQLA_TYPE_MAPPING = ModelConverter.SQLA_TYPE_MAPPING | {
        types.Enum: EnumField,
    }