Coverage for sqlmodel/_compat.py: 93%
290 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-08-12 19:13 +0000
1import types 1bacdef
2from contextlib import contextmanager 1bacdef
3from contextvars import ContextVar 1bacdef
4from dataclasses import dataclass 1bacdef
5from typing import ( 1bacdef
6 TYPE_CHECKING,
7 AbstractSet,
8 Any,
9 Callable,
10 Dict,
11 ForwardRef,
12 Generator,
13 Mapping,
14 Optional,
15 Set,
16 Type,
17 TypeVar,
18 Union,
19)
21from pydantic import VERSION as P_VERSION 1bacdef
22from pydantic import BaseModel 1bacdef
23from pydantic.fields import FieldInfo 1bacdef
24from typing_extensions import get_args, get_origin 1bacdef
26# Reassign variable to make it reexported for mypy
27PYDANTIC_VERSION = P_VERSION 1bacdef
28IS_PYDANTIC_V2 = PYDANTIC_VERSION.startswith("2.") 1bacdef
31if TYPE_CHECKING: 1bacdef
32 from .main import RelationshipInfo, SQLModel
34UnionType = getattr(types, "UnionType", Union) 1bacdef
35NoneType = type(None) 1bacdef
36T = TypeVar("T") 1bacdef
37InstanceOrType = Union[T, Type[T]] 1bacdef
38_TSQLModel = TypeVar("_TSQLModel", bound="SQLModel") 1bacdef
41class FakeMetadata: 1bacdef
42 max_length: Optional[int] = None 1bacdef
43 max_digits: Optional[int] = None 1bacdef
44 decimal_places: Optional[int] = None 1bacdef
47@dataclass 1bacdef
48class ObjectWithUpdateWrapper: 1acdef
49 obj: Any 1bacdef
50 update: Dict[str, Any] 1bacdef
52 def __getattribute__(self, __name: str) -> Any: 1bacdef
53 update = super().__getattribute__("update") 1ba
54 obj = super().__getattribute__("obj") 1ba
55 if __name in update: 1ba
56 return update[__name] 1ba
57 return getattr(obj, __name) 1ba
60def _is_union_type(t: Any) -> bool: 1bacdef
61 return t is UnionType or t is Union 1ba
64finish_init: ContextVar[bool] = ContextVar("finish_init", default=True) 1bacdef
67@contextmanager 1bacdef
68def partial_init() -> Generator[None, None, None]: 1bacdef
69 token = finish_init.set(False) 1ba
70 yield 1ba
71 finish_init.reset(token) 1ba
74if IS_PYDANTIC_V2: 1bacdef
75 from annotated_types import MaxLen 1ba
76 from pydantic import ConfigDict as BaseConfig 1ba
77 from pydantic._internal._fields import PydanticMetadata 1ba
78 from pydantic._internal._model_construction import ModelMetaclass 1ba
79 from pydantic._internal._repr import Representation as Representation 1ba
80 from pydantic_core import PydanticUndefined as Undefined 1ba
81 from pydantic_core import PydanticUndefinedType as UndefinedType 1ba
83 # Dummy for types, to make it importable
84 class ModelField: 1ba
85 pass 1ba
87 class SQLModelConfig(BaseConfig, total=False): 1ba
88 table: Optional[bool] 1ba
89 registry: Optional[Any] 1ba
91 def get_config_value( 1a
92 *, model: InstanceOrType["SQLModel"], parameter: str, default: Any = None
93 ) -> Any:
94 return model.model_config.get(parameter, default) 1ba
96 def set_config_value( 1a
97 *,
98 model: InstanceOrType["SQLModel"],
99 parameter: str,
100 value: Any,
101 ) -> None:
102 model.model_config[parameter] = value # type: ignore[literal-required] 1ba
104 def get_model_fields(model: InstanceOrType[BaseModel]) -> Dict[str, "FieldInfo"]: 1ba
105 return model.model_fields 1ba
107 def get_fields_set( 1a
108 object: InstanceOrType["SQLModel"],
109 ) -> Union[Set[str], Callable[[BaseModel], Set[str]]]:
110 return object.model_fields_set 1ba
112 def init_pydantic_private_attrs(new_object: InstanceOrType["SQLModel"]) -> None: 1ba
113 object.__setattr__(new_object, "__pydantic_fields_set__", set()) 1ba
114 object.__setattr__(new_object, "__pydantic_extra__", None) 1ba
115 object.__setattr__(new_object, "__pydantic_private__", None) 1ba
117 def get_annotations(class_dict: Dict[str, Any]) -> Dict[str, Any]: 1ba
118 return class_dict.get("__annotations__", {}) 1ba
120 def is_table_model_class(cls: Type[Any]) -> bool: 1ba
121 config = getattr(cls, "model_config", {}) 1ba
122 if config: 1ba
123 return config.get("table", False) or False 1ba
124 return False 1ba
126 def get_relationship_to( 1a
127 name: str,
128 rel_info: "RelationshipInfo",
129 annotation: Any,
130 ) -> Any:
131 origin = get_origin(annotation) 1ba
132 use_annotation = annotation 1ba
133 # Direct relationships (e.g. 'Team' or Team) have None as an origin
134 if origin is None: 1ba
135 if isinstance(use_annotation, ForwardRef): 1ba
136 use_annotation = use_annotation.__forward_arg__ 1ba
137 else:
138 return use_annotation 1ba
139 # If Union (e.g. Optional), get the real field
140 elif _is_union_type(origin): 1ba
141 use_annotation = get_args(annotation) 1ba
142 if len(use_annotation) > 2: 1ba
143 raise ValueError(
144 "Cannot have a (non-optional) union as a SQLAlchemy field"
145 )
146 arg1, arg2 = use_annotation 1ba
147 if arg1 is NoneType and arg2 is not NoneType: 1ba
148 use_annotation = arg2
149 elif arg2 is NoneType and arg1 is not NoneType: 1ba
150 use_annotation = arg1 1ba
151 else:
152 raise ValueError(
153 "Cannot have a Union of None and None as a SQLAlchemy field"
154 )
156 # If a list, then also get the real field
157 elif origin is list: 1ba
158 use_annotation = get_args(annotation)[0] 1ba
160 return get_relationship_to( 1ba
161 name=name, rel_info=rel_info, annotation=use_annotation
162 )
164 def is_field_noneable(field: "FieldInfo") -> bool: 1ba
165 if getattr(field, "nullable", Undefined) is not Undefined: 1ba
166 return field.nullable # type: ignore 1ba
167 origin = get_origin(field.annotation) 1ba
168 if origin is not None and _is_union_type(origin): 1ba
169 args = get_args(field.annotation) 1ba
170 if any(arg is NoneType for arg in args): 1ba
171 return True 1ba
172 if not field.is_required(): 1ba
173 if field.default is Undefined: 1ba
174 return False
175 if field.annotation is None or field.annotation is NoneType: # type: ignore[comparison-overlap] 1ba
176 return True
177 return False 1ba
178 return False 1ba
180 def get_type_from_field(field: Any) -> Any: 1ba
181 type_: Any = field.annotation 1ba
182 # Resolve Optional fields
183 if type_ is None: 1ba
184 raise ValueError("Missing field type")
185 origin = get_origin(type_) 1ba
186 if origin is None: 1ba
187 return type_ 1ba
188 if _is_union_type(origin): 1ba
189 bases = get_args(type_) 1ba
190 if len(bases) > 2: 1ba
191 raise ValueError(
192 "Cannot have a (non-optional) union as a SQLAlchemy field"
193 )
194 # Non optional unions are not allowed
195 if bases[0] is not NoneType and bases[1] is not NoneType: 1ba
196 raise ValueError( 1ba
197 "Cannot have a (non-optional) union as a SQLAlchemy field"
198 )
199 # Optional unions are allowed
200 return bases[0] if bases[0] is not NoneType else bases[1] 1ba
201 return origin 1ba
203 def get_field_metadata(field: Any) -> Any: 1ba
204 for meta in field.metadata: 1ba
205 if isinstance(meta, (PydanticMetadata, MaxLen)): 1ba
206 return meta 1ba
207 return FakeMetadata() 1ba
209 def post_init_field_info(field_info: FieldInfo) -> None: 1ba
210 return None 1ba
212 # Dummy to make it importable
213 def _calculate_keys( 1a
214 self: "SQLModel", 1a
215 include: Optional[Mapping[Union[int, str], Any]], 1a
216 exclude: Optional[Mapping[Union[int, str], Any]], 1a
217 exclude_unset: bool, 1a
218 update: Optional[Dict[str, Any]] = None, 1ba
219 ) -> Optional[AbstractSet[str]]: # pragma: no cover 1ba
220 return None
222 def sqlmodel_table_construct( 1a
223 *,
224 self_instance: _TSQLModel,
225 values: Dict[str, Any],
226 _fields_set: Union[Set[str], None] = None,
227 ) -> _TSQLModel:
228 # Copy from Pydantic's BaseModel.construct()
229 # Ref: https://github.com/pydantic/pydantic/blob/v2.5.2/pydantic/main.py#L198
230 # Modified to not include everything, only the model fields, and to
231 # set relationships
232 # SQLModel override to get class SQLAlchemy __dict__ attributes and
233 # set them back in after creating the object
234 # new_obj = cls.__new__(cls)
235 cls = type(self_instance) 1ba
236 old_dict = self_instance.__dict__.copy() 1ba
237 # End SQLModel override
239 fields_values: Dict[str, Any] = {} 1ba
240 defaults: Dict[ 1a
241 str, Any
242 ] = {} # keeping this separate from `fields_values` helps us compute `_fields_set`
243 for name, field in cls.model_fields.items(): 1ba
244 if field.alias and field.alias in values: 1ba
245 fields_values[name] = values.pop(field.alias)
246 elif name in values: 1ba
247 fields_values[name] = values.pop(name) 1ba
248 elif not field.is_required(): 1ba
249 defaults[name] = field.get_default(call_default_factory=True) 1ba
250 if _fields_set is None: 1ba
251 _fields_set = set(fields_values.keys()) 1ba
252 fields_values.update(defaults) 1ba
254 _extra: Union[Dict[str, Any], None] = None 1ba
255 if cls.model_config.get("extra") == "allow": 1ba
256 _extra = {}
257 for k, v in values.items():
258 _extra[k] = v
259 # SQLModel override, do not include everything, only the model fields
260 # else:
261 # fields_values.update(values)
262 # End SQLModel override
263 # SQLModel override
264 # Do not set __dict__, instead use setattr to trigger SQLAlchemy
265 # object.__setattr__(new_obj, "__dict__", fields_values)
266 # instrumentation
267 for key, value in {**old_dict, **fields_values}.items(): 1ba
268 setattr(self_instance, key, value) 1ba
269 # End SQLModel override
270 object.__setattr__(self_instance, "__pydantic_fields_set__", _fields_set) 1ba
271 if not cls.__pydantic_root_model__: 1ba
272 object.__setattr__(self_instance, "__pydantic_extra__", _extra) 1ba
274 if cls.__pydantic_post_init__: 1ba
275 self_instance.model_post_init(None)
276 elif not cls.__pydantic_root_model__: 1ba
277 # Note: if there are any private attributes, cls.__pydantic_post_init__ would exist
278 # Since it doesn't, that means that `__pydantic_private__` should be set to None
279 object.__setattr__(self_instance, "__pydantic_private__", None) 1ba
280 # SQLModel override, set relationships
281 # Get and set any relationship objects
282 for key in self_instance.__sqlmodel_relationships__: 1ba
283 value = values.get(key, Undefined) 1ba
284 if value is not Undefined: 1ba
285 setattr(self_instance, key, value) 1ba
286 # End SQLModel override
287 return self_instance 1ba
289 def sqlmodel_validate( 1a
290 cls: Type[_TSQLModel],
291 obj: Any,
292 *,
293 strict: Union[bool, None] = None,
294 from_attributes: Union[bool, None] = None,
295 context: Union[Dict[str, Any], None] = None,
296 update: Union[Dict[str, Any], None] = None,
297 ) -> _TSQLModel:
298 if not is_table_model_class(cls): 1ba
299 new_obj: _TSQLModel = cls.__new__(cls) 1ba
300 else:
301 # If table, create the new instance normally to make SQLAlchemy create
302 # the _sa_instance_state attribute
303 # The wrapper of this function should use with _partial_init()
304 with partial_init(): 1ba
305 new_obj = cls() 1ba
306 # SQLModel Override to get class SQLAlchemy __dict__ attributes and
307 # set them back in after creating the object
308 old_dict = new_obj.__dict__.copy() 1ba
309 use_obj = obj 1ba
310 if isinstance(obj, dict) and update: 1ba
311 use_obj = {**obj, **update}
312 elif update: 1ba
313 use_obj = ObjectWithUpdateWrapper(obj=obj, update=update) 1ba
314 cls.__pydantic_validator__.validate_python( 1ba
315 use_obj,
316 strict=strict,
317 from_attributes=from_attributes,
318 context=context,
319 self_instance=new_obj,
320 )
321 # Capture fields set to restore it later
322 fields_set = new_obj.__pydantic_fields_set__.copy() 1ba
323 if not is_table_model_class(cls): 1ba
324 # If not table, normal Pydantic code, set __dict__
325 new_obj.__dict__ = {**old_dict, **new_obj.__dict__} 1ba
326 else:
327 # Do not set __dict__, instead use setattr to trigger SQLAlchemy
328 # instrumentation
329 for key, value in {**old_dict, **new_obj.__dict__}.items(): 1ba
330 setattr(new_obj, key, value) 1ba
331 # Restore fields set
332 object.__setattr__(new_obj, "__pydantic_fields_set__", fields_set) 1ba
333 # Get and set any relationship objects
334 if is_table_model_class(cls): 1ba
335 for key in new_obj.__sqlmodel_relationships__: 1ba
336 value = getattr(use_obj, key, Undefined) 1ba
337 if value is not Undefined: 1ba
338 setattr(new_obj, key, value)
339 return new_obj 1ba
341 def sqlmodel_init(*, self: "SQLModel", data: Dict[str, Any]) -> None: 1ba
342 old_dict = self.__dict__.copy() 1ba
343 if not is_table_model_class(self.__class__): 1ba
344 self.__pydantic_validator__.validate_python( 1ba
345 data,
346 self_instance=self,
347 )
348 else:
349 sqlmodel_table_construct( 1ba
350 self_instance=self,
351 values=data,
352 )
353 object.__setattr__( 1ba
354 self,
355 "__dict__",
356 {**old_dict, **self.__dict__},
357 )
359else:
360 from pydantic import BaseConfig as BaseConfig # type: ignore[assignment] 1cdef
361 from pydantic.errors import ConfigError 1cdef
362 from pydantic.fields import ( # type: ignore[attr-defined, no-redef] 1cdef
363 SHAPE_SINGLETON,
364 ModelField,
365 )
366 from pydantic.fields import ( # type: ignore[attr-defined, no-redef] 1cdef
367 Undefined as Undefined, # noqa
368 )
369 from pydantic.fields import ( # type: ignore[attr-defined, no-redef] 1cdef
370 UndefinedType as UndefinedType,
371 )
372 from pydantic.main import ( # type: ignore[no-redef] 1cdef
373 ModelMetaclass as ModelMetaclass,
374 )
375 from pydantic.main import validate_model 1cdef
376 from pydantic.typing import resolve_annotations 1cdef
377 from pydantic.utils import ROOT_KEY, ValueItems 1cdef
378 from pydantic.utils import ( # type: ignore[no-redef] 1cdef
379 Representation as Representation,
380 )
382 class SQLModelConfig(BaseConfig): # type: ignore[no-redef] 1cdef
383 table: Optional[bool] = None # type: ignore[misc] 1cdef
384 registry: Optional[Any] = None # type: ignore[misc] 1cdef
386 def get_config_value( 1cdef
387 *, model: InstanceOrType["SQLModel"], parameter: str, default: Any = None
388 ) -> Any:
389 return getattr(model.__config__, parameter, default) # type: ignore[union-attr] 1cdef
391 def set_config_value( 1cdef
392 *,
393 model: InstanceOrType["SQLModel"],
394 parameter: str,
395 value: Any,
396 ) -> None:
397 setattr(model.__config__, parameter, value) # type: ignore 1cdef
399 def get_model_fields(model: InstanceOrType[BaseModel]) -> Dict[str, "FieldInfo"]: 1cdef
400 return model.__fields__ # type: ignore 1cdef
402 def get_fields_set( 1cdef
403 object: InstanceOrType["SQLModel"],
404 ) -> Union[Set[str], Callable[[BaseModel], Set[str]]]:
405 return object.__fields_set__ 1cdef
407 def init_pydantic_private_attrs(new_object: InstanceOrType["SQLModel"]) -> None: 1cdef
408 object.__setattr__(new_object, "__fields_set__", set()) 1cdef
410 def get_annotations(class_dict: Dict[str, Any]) -> Dict[str, Any]: 1cdef
411 return resolve_annotations( # type: ignore[no-any-return] 1cdef
412 class_dict.get("__annotations__", {}),
413 class_dict.get("__module__", None),
414 )
416 def is_table_model_class(cls: Type[Any]) -> bool: 1cdef
417 config = getattr(cls, "__config__", None) 1cdef
418 if config: 1cdef
419 return getattr(config, "table", False) 1cdef
420 return False
422 def get_relationship_to( 1cdef
423 name: str,
424 rel_info: "RelationshipInfo",
425 annotation: Any,
426 ) -> Any:
427 temp_field = ModelField.infer( # type: ignore[attr-defined] 1cdef
428 name=name,
429 value=rel_info,
430 annotation=annotation,
431 class_validators=None,
432 config=SQLModelConfig,
433 )
434 relationship_to = temp_field.type_ 1cdef
435 if isinstance(temp_field.type_, ForwardRef): 1cdef
436 relationship_to = temp_field.type_.__forward_arg__ 1cdef
437 return relationship_to 1cdef
439 def is_field_noneable(field: "FieldInfo") -> bool: 1cdef
440 if not field.required: # type: ignore[attr-defined] 1cdef
441 # Taken from [Pydantic](https://github.com/samuelcolvin/pydantic/blob/v1.8.2/pydantic/fields.py#L946-L947)
442 return field.allow_none and ( # type: ignore[attr-defined] 1cdef
443 field.shape != SHAPE_SINGLETON or not field.sub_fields # type: ignore[attr-defined]
444 )
445 return field.allow_none # type: ignore[no-any-return, attr-defined] 1cdef
447 def get_type_from_field(field: Any) -> Any: 1cdef
448 if isinstance(field.type_, type) and field.shape == SHAPE_SINGLETON: 1cdef
449 return field.type_ 1cdef
450 raise ValueError(f"The field {field.name} has no matching SQLAlchemy type") 1cdef
452 def get_field_metadata(field: Any) -> Any: 1cdef
453 metadata = FakeMetadata() 1cdef
454 metadata.max_length = field.field_info.max_length 1cdef
455 metadata.max_digits = getattr(field.type_, "max_digits", None) 1cdef
456 metadata.decimal_places = getattr(field.type_, "decimal_places", None) 1cdef
457 return metadata 1cdef
459 def post_init_field_info(field_info: FieldInfo) -> None: 1cdef
460 field_info._validate() # type: ignore[attr-defined] 1cdef
462 def _calculate_keys( 1cdef
463 self: "SQLModel",
464 include: Optional[Mapping[Union[int, str], Any]],
465 exclude: Optional[Mapping[Union[int, str], Any]],
466 exclude_unset: bool,
467 update: Optional[Dict[str, Any]] = None,
468 ) -> Optional[AbstractSet[str]]:
469 if include is None and exclude is None and not exclude_unset: 1cdef
470 # Original in Pydantic:
471 # return None
472 # Updated to not return SQLAlchemy attributes
473 # Do not include relationships as that would easily lead to infinite
474 # recursion, or traversing the whole database
475 return ( 1cdef
476 self.__fields__.keys() # noqa
477 ) # | self.__sqlmodel_relationships__.keys()
479 keys: AbstractSet[str]
480 if exclude_unset: 1cdef
481 keys = self.__fields_set__.copy() # noqa 1cdef
482 else:
483 # Original in Pydantic:
484 # keys = self.__dict__.keys()
485 # Updated to not return SQLAlchemy attributes
486 # Do not include relationships as that would easily lead to infinite
487 # recursion, or traversing the whole database
488 keys = (
489 self.__fields__.keys() # noqa
490 ) # | self.__sqlmodel_relationships__.keys()
491 if include is not None: 1cdef
492 keys &= include.keys()
494 if update: 1cdef
495 keys -= update.keys()
497 if exclude: 1cdef
498 keys -= {k for k, v in exclude.items() if ValueItems.is_true(v)}
500 return keys 1cdef
502 def sqlmodel_validate( 1cdef
503 cls: Type[_TSQLModel],
504 obj: Any,
505 *,
506 strict: Union[bool, None] = None,
507 from_attributes: Union[bool, None] = None,
508 context: Union[Dict[str, Any], None] = None,
509 update: Union[Dict[str, Any], None] = None,
510 ) -> _TSQLModel:
511 # This was SQLModel's original from_orm() for Pydantic v1
512 # Duplicated from Pydantic
513 if not cls.__config__.orm_mode: # type: ignore[attr-defined] # noqa 1cdef
514 raise ConfigError(
515 "You must have the config attribute orm_mode=True to use from_orm"
516 )
517 if not isinstance(obj, Mapping): 1cdef
518 obj = ( 1cdef
519 {ROOT_KEY: obj}
520 if cls.__custom_root_type__ # type: ignore[attr-defined] # noqa
521 else cls._decompose_class(obj) # type: ignore[attr-defined] # noqa
522 )
523 # SQLModel, support update dict
524 if update is not None: 1cdef
525 obj = {**obj, **update} 1cdef
526 # End SQLModel support dict
527 if not getattr(cls.__config__, "table", False): # noqa 1cdef
528 # If not table, normal Pydantic code
529 m: _TSQLModel = cls.__new__(cls) 1cdef
530 else:
531 # If table, create the new instance normally to make SQLAlchemy create
532 # the _sa_instance_state attribute
533 m = cls() 1cdef
534 values, fields_set, validation_error = validate_model(cls, obj) 1cdef
535 if validation_error: 1cdef
536 raise validation_error
537 # Updated to trigger SQLAlchemy internal handling
538 if not getattr(cls.__config__, "table", False): # noqa 1cdef
539 object.__setattr__(m, "__dict__", values) 1cdef
540 else:
541 for key, value in values.items(): 1cdef
542 setattr(m, key, value) 1cdef
543 # Continue with standard Pydantic logic
544 object.__setattr__(m, "__fields_set__", fields_set) 1cdef
545 m._init_private_attributes() # type: ignore[attr-defined] # noqa 1cdef
546 return m 1cdef
548 def sqlmodel_init(*, self: "SQLModel", data: Dict[str, Any]) -> None: 1cdef
549 values, fields_set, validation_error = validate_model(self.__class__, data) 1cdef
550 # Only raise errors if not a SQLModel model
551 if ( 1cde
552 not is_table_model_class(self.__class__) # noqa
553 and validation_error
554 ):
555 raise validation_error 1cdef
556 if not is_table_model_class(self.__class__): 1cdef
557 object.__setattr__(self, "__dict__", values) 1cdef
558 else:
559 # Do not set values as in Pydantic, pass them through setattr, so
560 # SQLAlchemy can handle them
561 for key, value in values.items(): 1cdef
562 setattr(self, key, value) 1cdef
563 object.__setattr__(self, "__fields_set__", fields_set) 1cdef
564 non_pydantic_keys = data.keys() - values.keys() 1cdef
566 if is_table_model_class(self.__class__): 1cdef
567 for key in non_pydantic_keys: 1cdef
568 if key in self.__sqlmodel_relationships__: 1cdef
569 setattr(self, key, data[key]) 1cdef