def get_model_mapper(model, stoppage=None, full=True):
"""Get a dictionary of name: class for all the objects in model."""
model = get_model(model)
flat_models = get_flat_models_from_model(model)
# this is the list of all the referenced objects
model_name_map = get_model_name_map(flat_models)
# flip the dictionary so I can access each class by name
model_name_map = {v: k for k, v in model_name_map.items()}
if full:
if not stoppage:
stoppage = set(
['NoExtraBaseModel', 'ModelMetaclass', 'BaseModel', 'object'])
# Pydantic does not necessarily add all the baseclasses to the OpenAPI
# documentation. We check all of them and them to the list if they are not
# already added
models = list(model_name_map.values())
for model in models:
for cls in type.mro(model):
if cls.__name__ in stoppage:
break
if cls.__name__ not in model_name_map:
model_name_map[cls.__name__] = cls
return model_name_map
def test_flat_models_with_submodels():
class Foo(BaseModel):
a: str
class Bar(BaseModel):
b: List[Foo]
class Baz(BaseModel):
c: Dict[str, Bar]
flat_models = get_flat_models_from_model(Baz)
assert flat_models == set([Foo, Bar, Baz])
def api_sanitize(
pydantic_model: Type[BaseModel],
fields_to_leave: Optional[List[str]] = None,
allow_dict_msonable=False,
):
"""
Function to clean up pydantic models for the API by:
1.) Making fields optional
2.) Allowing dictionaries in-place of the objects for MSONable quantities
WARNING: This works in place, so it mutates the model and all sub-models
Args:
fields_to_leave: list of strings for model fields as "model__name__.field"
"""
models = [
model
for model in get_flat_models_from_model(pydantic_model)
if issubclass(model, BaseModel)
] # type: List[Type[BaseModel]]
fields_to_leave = fields_to_leave or []
fields_tuples = [f.split(".") for f in fields_to_leave]
assert all(len(f) == 2 for f in fields_tuples)
for model in models:
model_fields_to_leave = {f[1] for f in fields_tuples if model.__name__ == f[0]}
for name, field in model.__fields__.items():
field_type = field.type_
if name not in model_fields_to_leave:
field.required = False
field.field_info.default = None
if field_type is not None and allow_dict_msonable:
if lenient_issubclass(field_type, MSONable):
field.type_ = allow_msonable_dict(field_type)
else:
for sub_type in get_args(field_type):
if lenient_issubclass(sub_type, MSONable):
allow_msonable_dict(sub_type)
field.populate_validators()
return pydantic_model
