53 lines
1.6 KiB
Python

from typing import TYPE_CHECKING, Any, Dict, Optional, Type
from pydantic_core import core_schema
from typing_extensions import Annotated
from ninja import Body
from ninja.utils import is_optional_type
class ModelToDict(dict):
_wrapped_model: Any = None
_wrapped_model_dump_params: Dict[str, Any] = {}
@classmethod
def __get_pydantic_core_schema__(cls, _source: Any, _handler: Any) -> Any:
return core_schema.no_info_after_validator_function(
cls._validate,
cls._wrapped_model.__pydantic_core_schema__,
)
@classmethod
def _validate(cls, input_value: Any) -> Any:
return input_value.model_dump(**cls._wrapped_model_dump_params)
def create_patch_schema(schema_cls: Type[Any]) -> Type[ModelToDict]:
values, annotations = {}, {}
for f in schema_cls.__fields__.keys():
t = schema_cls.__annotations__[f]
if not is_optional_type(t):
values[f] = getattr(schema_cls, f, None)
annotations[f] = Optional[t]
values["__annotations__"] = annotations
OptionalSchema = type(f"{schema_cls.__name__}Patch", (schema_cls,), values)
class OptionalDictSchema(ModelToDict):
_wrapped_model = OptionalSchema
_wrapped_model_dump_params = {"exclude_unset": True}
return OptionalDictSchema
class PatchDictUtil:
def __getitem__(self, schema_cls: Any) -> Any:
new_cls = create_patch_schema(schema_cls)
return Body[new_cls] # type: ignore
if TYPE_CHECKING: # pragma: nocover
PatchDict = Annotated[dict, "<PatchDict>"]
else:
PatchDict = PatchDictUtil()