Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add union type and subtypes check in schema model signature #1242

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+45 −16
Open

Add union type and subtypes check in schema model signature #1242

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 45 additions & 16 deletions ninja/signature/details.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,13 +198,20 @@ def _args_flatten_map(self, args: List[FuncParam]) -> Dict[str, Tuple[str, ...]]

def _model_flatten_map(self, model: TModel, prefix: str) -> Generator:
field: FieldInfo
for attr, field in model.model_fields.items():
field_name = field.alias or attr
name = f"{prefix}{self.FLATTEN_PATH_SEP}{field_name}"
if is_pydantic_model(field.annotation):
yield from self._model_flatten_map(field.annotation, name) # type: ignore
else:
yield field_name, name
if get_origin(model) in UNION_TYPES:
# If the model is a union type, process each type in the union
for arg in get_args(model):
if type(arg) is None:
continue # Skip NoneType
yield from self._model_flatten_map(arg, prefix)
else:
for attr, field in model.model_fields.items():
field_name = field.alias or attr
name = f"{prefix}{self.FLATTEN_PATH_SEP}{field_name}"
if is_pydantic_model(field.annotation):
yield from self._model_flatten_map(field.annotation, name) # type: ignore
else:
yield field_name, name

def _get_param_type(self, name: str, arg: inspect.Parameter) -> FuncParam:
# _EMPTY = self.signature.empty
Expand Down Expand Up @@ -278,7 +285,11 @@ def _get_param_type(self, name: str, arg: inspect.Parameter) -> FuncParam:
def is_pydantic_model(cls: Any) -> bool:
try:
if get_origin(cls) in UNION_TYPES:
return any(issubclass(arg, pydantic.BaseModel) for arg in get_args(cls))
return any(
issubclass(arg, pydantic.BaseModel)
for arg in get_args(cls)
if (type(arg) is not None)
)
return issubclass(cls, pydantic.BaseModel)
except TypeError:
return False
Expand Down Expand Up @@ -321,14 +332,32 @@ def detect_collection_fields(
for attr in path[1:]:
if hasattr(annotation_or_field, "annotation"):
annotation_or_field = annotation_or_field.annotation
annotation_or_field = next(
(
a
for a in annotation_or_field.model_fields.values()
if a.alias == attr
),
annotation_or_field.model_fields.get(attr),
) # pragma: no cover

# check union types
if get_origin(annotation_or_field) in UNION_TYPES:
for arg in get_args(annotation_or_field):
if type(arg) is None:
continue # Skip NoneType
if hasattr(arg, "model_fields"):
annotation_or_field = next(
(
a
for a in arg.model_fields.values()
if a.alias == attr
),
arg.model_fields.get(attr),
) # pragma: no cover
else:
continue
else:
annotation_or_field = next(
(
a
for a in annotation_or_field.model_fields.values()
if a.alias == attr
),
annotation_or_field.model_fields.get(attr),
) # pragma: no cover

annotation_or_field = getattr(
annotation_or_field, "outer_type_", annotation_or_field
Expand Down
Loading