This commit is contained in:
Diogo Sousa
2024-03-20 14:54:14 +00:00
parent b932b52d05
commit 62ca11d687
2 changed files with 25 additions and 23 deletions

View File

@@ -83,30 +83,32 @@ class OBBject(Tagged, Generic[T]):
@classmethod
def results_type_repr(cls, params: Optional[Any] = None) -> str:
"""Return the results type name."""
"""Return the results type representation."""
results_field = cls.model_fields.get("results")
type_ = params[0] if params else results_field.annotation
name = type_.__name__ if hasattr(type_, "__name__") else str(type_)
type_repr = "Any"
if results_field:
type_ = params[0] if params else results_field.annotation
type_repr = getattr(type_, "__name__", str(type_))
if (json_schema_extra := results_field.json_schema_extra) is not None:
model = json_schema_extra.get("model")
if json_schema_extra := getattr(results_field, "json_schema_extra", {}):
model = json_schema_extra.get("model", "Any")
if json_schema_extra.get("is_union"):
return f"Union[List[{model}], {model}]"
if json_schema_extra.get("has_list"):
return f"List[{model}]"
if json_schema_extra.get("is_union"):
return f"Union[List[{model}], {model}]"
if json_schema_extra.get("has_list"):
return f"List[{model}]"
return model
return model
if "typing." in str(type_):
unpack_optional = sub(r"Optional\[(.*)\]", r"\1", str(type_))
name = sub(
r"(\w+\.)*(\w+)?(\, NoneType)?",
r"\2",
unpack_optional,
)
if "typing." in str(type_):
unpack_optional = sub(r"Optional\[(.*)\]", r"\1", str(type_))
type_repr = sub(
r"(\w+\.)*(\w+)?(\, NoneType)?",
r"\2",
unpack_optional,
)
return name
return type_repr
@classmethod
def model_parametrized_name(cls, params: Any) -> str:

View File

@@ -25,11 +25,11 @@ df = pd.DataFrame(
df_multiindex = df.set_index(["x", "y"])
simple_base_model = [
Data(x=i, y=j, z=k) for i in range(2) for j in range(6, 8) for k in range(10, 12)
Data(x=i, y=j, z=k) for i in range(2) for j in range(6, 8) for k in range(10, 12) # type: ignore[call-arg]
]
multi_index_base_model = [
Data(x=i, y=j, z=k, is_multiindex=True, multiindex_names="['x','y']")
Data(x=i, y=j, z=k, is_multiindex=True, multiindex_names="['x','y']") # type: ignore[call-arg]
for i in range(2)
for j in range(6, 8)
for k in range(10, 12)
@@ -93,11 +93,11 @@ def test_list_to_basemodel(data_list, expected):
"data_dict, expected",
[
# Simple dictionary
({"a": 10}, Data(a=10)),
({"a": 10}, Data(a=10)), # type: ignore[call-arg]
# Nested dictionary (assuming Data can handle nested dicts)
({"b": {"c": 20}}, Data(b={"c": 20})),
({"b": {"c": 20}}, Data(b={"c": 20})), # type: ignore[call-arg]
# Dictionary with list (assuming Data can handle lists)
({"d": [30, 40]}, Data(d=[30, 40])),
({"d": [30, 40]}, Data(d=[30, 40])), # type: ignore[call-arg]
],
)
def test_dict_to_basemodel(data_dict, expected):