|
5 | 5 | import functools |
6 | 6 | import logging |
7 | 7 | from os import environ, path |
8 | | -from typing import Any, List, Sequence, Union |
| 8 | +from typing import Any, List, Sequence, Union, get_args, get_origin |
9 | 9 | import warnings |
10 | 10 |
|
11 | 11 | from google.protobuf.text_format import MessageToString |
@@ -170,25 +170,22 @@ def stringify_model(model: BaseModel) -> Sequence[str]: |
170 | 170 | It is assumed optional field does not contain other models (only basic types). |
171 | 171 | """ |
172 | 172 |
|
173 | | - def _lookup_type(current_model: BaseModel, type_path: str) -> str: |
174 | | - splitted_path = type_path.split(".", maxsplit=1) |
175 | | - value = current_model.__annotations__[splitted_path[0]] |
176 | | - if len(splitted_path) == 1: |
177 | | - if value.__dict__.get("__origin__") is Union: |
178 | | - return ", ".join(arg.__name__ for arg in value.__args__) |
179 | | - return value.__name__ |
180 | | - return _lookup_type(value, splitted_path[1]) |
| 173 | + def _lookup_type(current_model: BaseModel, field_name: str) -> str: |
| 174 | + value = type(current_model).model_fields[field_name].annotation |
| 175 | + if value is None: |
| 176 | + return "Any" |
| 177 | + if get_origin(value) is Union: |
| 178 | + return ", ".join(arg.__name__ for arg in get_args(value)) |
| 179 | + return value.__name__ |
181 | 180 |
|
182 | 181 | def _recurse_into( |
183 | 182 | current_model: BaseModel, prefix: str, output: List[str] |
184 | 183 | ) -> Sequence[str]: |
185 | 184 | for name, field in dict(current_model).items(): |
186 | | - if hasattr(field, "__annotations__"): |
187 | | - _recurse_into( |
188 | | - getattr(current_model, name), (prefix or "") + f"{name}.", output |
189 | | - ) |
| 185 | + if isinstance(field, BaseModel): |
| 186 | + _recurse_into(field, (prefix or "") + f"{name}.", output) |
190 | 187 | else: |
191 | | - field_type = _lookup_type(model, f"{prefix}{name}") |
| 188 | + field_type = _lookup_type(current_model, name) |
192 | 189 | output.append(f"{prefix}{name} = {field} ({field_type})") |
193 | 190 | return output |
194 | 191 |
|
|
0 commit comments