root_model.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. """RootModel class and type definitions."""
  2. from __future__ import annotations as _annotations
  3. from copy import copy, deepcopy
  4. from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar
  5. from pydantic_core import PydanticUndefined
  6. from typing_extensions import Self, dataclass_transform
  7. from . import PydanticUserError
  8. from ._internal import _model_construction, _repr
  9. from .main import BaseModel, _object_setattr
  10. if TYPE_CHECKING:
  11. from .fields import Field as PydanticModelField
  12. from .fields import PrivateAttr as PydanticModelPrivateAttr
  13. # dataclass_transform could be applied to RootModel directly, but `ModelMetaclass`'s dataclass_transform
  14. # takes priority (at least with pyright). We trick type checkers into thinking we apply dataclass_transform
  15. # on a new metaclass.
  16. @dataclass_transform(kw_only_default=False, field_specifiers=(PydanticModelField, PydanticModelPrivateAttr))
  17. class _RootModelMetaclass(_model_construction.ModelMetaclass): ...
  18. else:
  19. _RootModelMetaclass = _model_construction.ModelMetaclass
  20. __all__ = ('RootModel',)
  21. RootModelRootType = TypeVar('RootModelRootType')
  22. class RootModel(BaseModel, Generic[RootModelRootType], metaclass=_RootModelMetaclass):
  23. """!!! abstract "Usage Documentation"
  24. [`RootModel` and Custom Root Types](../concepts/models.md#rootmodel-and-custom-root-types)
  25. A Pydantic `BaseModel` for the root object of the model.
  26. Attributes:
  27. root: The root object of the model.
  28. __pydantic_root_model__: Whether the model is a RootModel.
  29. __pydantic_private__: Private fields in the model.
  30. __pydantic_extra__: Extra fields in the model.
  31. """
  32. __pydantic_root_model__ = True
  33. __pydantic_private__ = None
  34. __pydantic_extra__ = None
  35. root: RootModelRootType
  36. def __init_subclass__(cls, **kwargs):
  37. extra = cls.model_config.get('extra')
  38. if extra is not None:
  39. raise PydanticUserError(
  40. "`RootModel` does not support setting `model_config['extra']`", code='root-model-extra'
  41. )
  42. super().__init_subclass__(**kwargs)
  43. def __init__(self, /, root: RootModelRootType = PydanticUndefined, **data) -> None: # type: ignore
  44. __tracebackhide__ = True
  45. if data:
  46. if root is not PydanticUndefined:
  47. raise ValueError(
  48. '"RootModel.__init__" accepts either a single positional argument or arbitrary keyword arguments'
  49. )
  50. root = data # type: ignore
  51. self.__pydantic_validator__.validate_python(root, self_instance=self)
  52. __init__.__pydantic_base_init__ = True # pyright: ignore[reportFunctionMemberAccess]
  53. @classmethod
  54. def model_construct(cls, root: RootModelRootType, _fields_set: set[str] | None = None) -> Self: # type: ignore
  55. """Create a new model using the provided root object and update fields set.
  56. Args:
  57. root: The root object of the model.
  58. _fields_set: The set of fields to be updated.
  59. Returns:
  60. The new model.
  61. Raises:
  62. NotImplemented: If the model is not a subclass of `RootModel`.
  63. """
  64. return super().model_construct(root=root, _fields_set=_fields_set)
  65. def __getstate__(self) -> dict[Any, Any]:
  66. return {
  67. '__dict__': self.__dict__,
  68. '__pydantic_fields_set__': self.__pydantic_fields_set__,
  69. }
  70. def __setstate__(self, state: dict[Any, Any]) -> None:
  71. _object_setattr(self, '__pydantic_fields_set__', state['__pydantic_fields_set__'])
  72. _object_setattr(self, '__dict__', state['__dict__'])
  73. def __copy__(self) -> Self:
  74. """Returns a shallow copy of the model."""
  75. cls = type(self)
  76. m = cls.__new__(cls)
  77. new_dict = copy(self.__dict__)
  78. new_dict['root'] = copy(self.__dict__['root'])
  79. _object_setattr(m, '__dict__', new_dict)
  80. _object_setattr(m, '__pydantic_fields_set__', copy(self.__pydantic_fields_set__))
  81. return m
  82. def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self:
  83. """Returns a deep copy of the model."""
  84. cls = type(self)
  85. m = cls.__new__(cls)
  86. _object_setattr(m, '__dict__', deepcopy(self.__dict__, memo=memo))
  87. # This next line doesn't need a deepcopy because __pydantic_fields_set__ is a set[str],
  88. # and attempting a deepcopy would be marginally slower.
  89. _object_setattr(m, '__pydantic_fields_set__', copy(self.__pydantic_fields_set__))
  90. return m
  91. if TYPE_CHECKING:
  92. def model_dump( # type: ignore
  93. self,
  94. *,
  95. mode: Literal['json', 'python'] | str = 'python',
  96. include: Any = None,
  97. exclude: Any = None,
  98. context: dict[str, Any] | None = None,
  99. by_alias: bool | None = None,
  100. exclude_unset: bool = False,
  101. exclude_defaults: bool = False,
  102. exclude_none: bool = False,
  103. exclude_computed_fields: bool = False,
  104. round_trip: bool = False,
  105. warnings: bool | Literal['none', 'warn', 'error'] = True,
  106. serialize_as_any: bool = False,
  107. ) -> Any:
  108. """This method is included just to get a more accurate return type for type checkers.
  109. It is included in this `if TYPE_CHECKING:` block since no override is actually necessary.
  110. See the documentation of `BaseModel.model_dump` for more details about the arguments.
  111. Generally, this method will have a return type of `RootModelRootType`, assuming that `RootModelRootType` is
  112. not a `BaseModel` subclass. If `RootModelRootType` is a `BaseModel` subclass, then the return
  113. type will likely be `dict[str, Any]`, as `model_dump` calls are recursive. The return type could
  114. even be something different, in the case of a custom serializer.
  115. Thus, `Any` is used here to catch all of these cases.
  116. """
  117. ...
  118. def __eq__(self, other: Any) -> bool:
  119. if not isinstance(other, RootModel):
  120. return NotImplemented
  121. return self.__pydantic_fields__['root'].annotation == other.__pydantic_fields__[
  122. 'root'
  123. ].annotation and super().__eq__(other)
  124. def __repr_args__(self) -> _repr.ReprArgs:
  125. yield 'root', self.root