52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
import pydantic
|
|
|
|
class BaseModel(pydantic.BaseModel):
|
|
|
|
def __init__(self, **data):
|
|
for key, value in data.items():
|
|
if isinstance(value, dict):
|
|
if 'type' in value:
|
|
data[key] = TypedModel.parse_obj(value)
|
|
super().__init__(**data)
|
|
|
|
# Adapted from https://github.com/pydantic/pydantic/discussions/3091
|
|
class TypedModel(BaseModel):
|
|
|
|
_subtypes_ = []
|
|
|
|
def __init_subclass__(cls, type=None):
|
|
cls._subtypes_.append([type, cls])
|
|
|
|
@classmethod
|
|
def get_cls(_cls, type):
|
|
for t, cls in _cls._subtypes_:
|
|
if t == type:
|
|
return cls
|
|
raise ValueError(f'Unknown type {type}')
|
|
|
|
@classmethod
|
|
def get_type(_cls, cls_name):
|
|
for t, cls in _cls._subtypes_:
|
|
if cls.__name__ == cls_name:
|
|
return t
|
|
raise ValueError(f'Unknown class {cls_name}')
|
|
|
|
@classmethod
|
|
def parse_obj(cls, obj):
|
|
data_type = obj.get('type')
|
|
if data_type is None:
|
|
raise ValueError(f'type is required for {cls.__name__}')
|
|
|
|
sub = cls.get_cls(data_type)
|
|
if sub is None:
|
|
raise ValueError(f'Unknown type {data_type}')
|
|
return sub(**obj)
|
|
|
|
def _iter(self, **kwargs):
|
|
yield 'type', self.get_type(self.__class__.__name__)
|
|
yield from super()._iter(**kwargs)
|
|
|
|
@property
|
|
def type(self):
|
|
return self.get_type(self.__class__.__name__)
|
|
|