Marshmallow Schema and Class Inheritance

Viewed 2712

I am using Marshmallow for the first time and unfortunately could not find an answer on the internet.

There are two classes, one inherits from the other. Both should be serialisable and deserialisable. After deserialisation they should be available again as Python objects. Therefore I use the post_load decorator. But it seems that this causes problems. In the minimal working example below I get the following exception: type object argument after ** must be a mapping, not Bicycle.

One could now look for the attribute amount_of_tires and, if necessary, pass on the data. But this does not feel like the right solution.

Is there a best practise to solve this problem?

from marshmallow import Schema, fields
from marshmallow.decorators import post_load


class Vehicle:
    def __init__(self, weight):
        self.weight = weight


class VehicleSchema(Schema):
    weight = fields.Float()

    @post_load
    def make_vehicle(self, data, **kwargs) -> Vehicle:
        return Vehicle(**data)


class Bicycle(Vehicle):
    def __init__(self, weight, amount_of_tires):
        Vehicle.__init__(self, weight=weight)
        self.amount_of_tires = amount_of_tires


class BicycleSchema(VehicleSchema):
    amount_of_tires = fields.Integer()

    @post_load
    def make_bicycle(self, data, **kwargs) -> Bicycle:
        return Bicycle(**data)


my_bicycle = Bicycle(weight=10, amount_of_tires=2)


schema = BicycleSchema()
json_string = schema.dumps(my_bicycle)

deserialised_my_bicycle = schema.loads(json_string)
print(deserialised_my_bicycle)


1 Answers

Both make_vehicle and make_bicycle are registered as post_load hooks, and so they are getting called one after the other. When make_vehicle is called, make_bicycle is already called, and so type(data) == Bicycle, hence the error you see.

I recommend the following solution:

class VehicleSchema(Schema):
    model_class = Vehicle

    weight = fields.Float()

    @post_load
    def make_vehicle(self, data, **kwargs) -> Vehicle:
        return type(self).model_class(**data)

class BicycleSchema(VehicleSchema):
    model_class = Bicycle
    amount_of_tires = fields.Integer()

To get more accurate type hints with this solution - one can use Generic and TypeVar on the base VehicleSchema class (see https://docs.python.org/3/library/typing.html#typing.Generic).

Related