Function flatten_pydantic_model

langroid/utils/pydantic_utils.py:104–154

Given a possibly nested Pydantic class, return a flattened version of it, by constructing top-level fields, whose names are formed from the path through the nested structure, separated by double underscores. Args: model (Type[BaseModel]): The Pydantic model to flatten.

(
    model: Type[BaseModel],
    base_model: Type[BaseModel] = BaseModel,
)

def flatten_pydantic_model(
    model: Type[BaseModel],
    base_model: Type[BaseModel] = BaseModel,
) -> Type[BaseModel]:
    """
    Given a possibly nested Pydantic class, return a flattened version of it,
    by constructing top-level fields, whose names are formed from the path
    through the nested structure, separated by double underscores.

    Args:
        model (Type[BaseModel]): The Pydantic model to flatten.
        base_model (Type[BaseModel], optional): The base model to use for the
            flattened model. Defaults to BaseModel.

    Returns:
        Type[BaseModel]: The flattened Pydantic model.
    """

    flattened_fields: Dict[str, Any] = {}
    models_to_process = [(model, "")]

    while models_to_process:
        current_model, current_prefix = models_to_process.pop()

        for name, field in current_model.model_fields.items():
            field_type = field.annotation if hasattr(field, "annotation") else field
            if isinstance(field_type, type) and issubclass(field_type, BaseModel):
                new_prefix = (
                    f"{current_prefix}{name}__" if current_prefix else f"{name}__"
                )
                models_to_process.append((field_type, new_prefix))
            else:
                flattened_name = f"{current_prefix}{name}"

                if (
                    hasattr(field, "default_factory")
                    and field.default_factory is not None
                ):
                    flattened_fields[flattened_name] = (
                        field_type,
                        field.default_factory,
                    )
                elif hasattr(field, "default") and field.default is not ...:
                    flattened_fields[flattened_name] = (
                        field_type,
                        field.default,
                    )
                else:
                    flattened_fields[flattened_name] = (field_type, ...)

    return create_model("FlatModel", __base__=base_model, **flattened_fields)


def get_field_names(model: Type[BaseModel]) -> List[str]: