Coverage for src / dynapydantic / annotations.py: 100%

1"""Custom annotations for dynapydantic""" 

2 

3import typing as ty 

4 

5from pydantic import GetCoreSchemaHandler, PydanticSchemaGenerationError 

6from pydantic_core import core_schema 

7 

8from .free_funcs import union 

9from .subclass_tracking_model import SubclassTrackingModel 

10from .tracking_group import TrackingGroup 

11 

12ModelT = ty.TypeVar("ModelT", bound=SubclassTrackingModel) 

13 

14 

15class PydanticAdapter: 

16 """Pydantic type adapter for SubclassTrackingModel""" 

17 

18 @staticmethod 

19 def __get_pydantic_core_schema__( 

20 source_type: type[SubclassTrackingModel], 

21 handler: GetCoreSchemaHandler, 

22 ) -> core_schema.CoreSchema: 

23 """Get the pydantic schema for this type""" 

24 return handler(union(source_type)) 

25 

26 

27if ty.TYPE_CHECKING: # pragma: no cover 

28 # From the type checker's perspective, Polymorphic[T] = Annotated[T] 

29 Polymorphic = ty.Annotated[ModelT, ...] 

30else: 

31 

32 class Polymorphic: 

33 """Annotation used to mark a type as having duck-typing behavior 

34 

35 This annotation is only valid for SubclassTrackingModel's. 

36 

37 Similar to SerializeAsAny, a field annotated with this shall serialize as 

38 according to its actual type, not the field annotation type. In addition, 

39 parsing will function as if the field annotation type were the union of 

40 all tracked subclasses. 

41 """ 

42 

43 def __class_getitem__(cls, item: ModelT) -> ty.Any: # noqa: ANN401 

44 """Get the annotation for the pydantic field""" 

45 if not isinstance(item, type): 

46 msg = f"dynapydantic.Polymorphic must be given a type, not {item}" 

47 raise TypeError(msg) 

48 

49 if getattr(item, "__DYNAPYDANTIC_IMPLICIT_POLYMORPHIC__", False): 

50 return item 

51 

52 if not issubclass(item, SubclassTrackingModel): 

53 msg = ( 

54 f"Polymorphic was given {item}, which was not a " 

55 "SubclassTrackingModel." 

56 ) 

57 raise PydanticSchemaGenerationError(msg) 

58 

59 return ty.Annotated[item, PydanticAdapter] 

60 

61 

62if ty.TYPE_CHECKING: # pragma: no cover 

63 

64 class Union: 

65 """Annotation used to get the union out of a dynapydantic entity""" 

66 

67 @ty.overload 

68 def __class_getitem__(cls, item: type[ModelT]) -> ModelT: ... 

69 

70 @ty.overload 

71 def __class_getitem__(cls, item: TrackingGroup) -> ty.Any: ... # noqa: ANN401 

72 

73 def __class_getitem__(cls, item: TrackingGroup | type[ModelT]) -> object: 

74 """Return the union""" 

75 

76else: 

77 

78 class Union: 

79 """Annotation used to get the union out of a dynapydantic entity 

80 

81 This annotation is primarily used for using the union of all models in 

82 a `TrackingGroup` as a field annotation. It can be used with 

83 `SubclassTrackingModel`, but in general, `Polymorphic` is preferable. 

84 """ 

85 

86 def __class_getitem__( 

87 cls, item: TrackingGroup | type[SubclassTrackingModel] 

88 ) -> ty.Any: # noqa: ANN401 

89 """Return the union""" 

90 return union(item)