Search code examples
pythonflaskflask-restplusflask-restx

Convert json field in body to enum with flask-restx


I have the following api definition for flask-restx (though should also work with flask-restplus). Is there someway to convert the enum-field in the request body the the Enum MyEnum without too much overhead or using DAOs?

class MyEnum(Enum):
    FOO = auto()
    BAR = auto()


@dataclass(frozen=True)
class MyClass:
    enum: MyEnum


api = Namespace('ns')

model = api.model('Model', {
    'enum': fields.String(enum=[x.name for x in MyEnum]),
})

@api.route('/')
class MyClass(Resource):

    @api.expect(Model)
    def post(self) -> None:
        c = MyClass(**api.payload)
        print(type(c.enum))  # <class 'str'> (but I want <enum 'MyEnum'>)
        assert(type(c.enum) == MyEnum)  # Fails

Solution

  • Ok I have written a decorator which will replace the enum value with the enum

    def decode_enum(api: Namespace, enum_cls: Type[Enum], keys: List[str]):
        def replace_item(obj: dict, keys_: List[str], new_value: Type[Enum]):
            if not keys_:
                return new_value
            obj[keys_[0]] = replace_item(obj[keys_[0]], keys_[1:], new_value)
            return obj
    
        def decoder(f):
            @wraps(f)
            def wrapper(*args, **kwds):
                value = api.payload
                for k in keys:
                    value = value[k]
                enum = enum_cls[value]
                api.payload[keys[0]] = replace_item(api.payload[keys[0]], keys[1:], enum)
                return f(*args, **kwds)
            return wrapper
        return decoder
    

    The usage would be like this

    @decode_enum(api, MyEnum, ['enum'])
    @api.expect(Model)
    def post(self) -> None:
        c = MyClass(**api.payload)
        print(type(c.enum))  # <enum 'MyEnum'>
    

    The replace_item function was inspired by this SO Answer: https://stackoverflow.com/a/45335542/6900162