CODEWITHSUNDEEP
pythonpydantic
WhatAmIDoing
WhatAmIDoing

Reputation: 220

How do I specify schema of an Enum field in a Pydantic Spark model

I have a Pydantic model which has an enum field. I'm trying to get the schema back but I get a runtime error which I suppose makes sense (I want a string type, but Pydantic is getting the Enum). How do I override this type?

>>> from pydantic_spark.base import SparkBase
>>> from enum import Enum

>>> class TestEnum(Enum):
>>>     SOMETHING = "something"
>>>     OTHER_THING = "something else"

>>> class TestModel(SparkBase):
>>>     thing: TestEnum

>>> my_model = TestModel(thing=TestEnum.SOMETHING)
>>> my_model.spark_schema()

RuntimeError: Unknown enum type: None

Upvotes: 0

Views: 369

Answers (1)

WhatAmIDoing
WhatAmIDoing

Reputation: 220

Edit: found the answer here.

This works fine:

class MyEnum(str, Enum):
    ...

Upvotes: 0

Related Questions