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

1"""Free function API for interacting with dynapydantic entities 

2 

3We are using free functions to avoid cluttering the namespace of 

4SubclassTrackingModel, following pydantic's lead: 

5https://github.com/pydantic/pydantic/issues/10032 

6""" 

7 

8import typing as ty 

9 

10import pydantic 

11 

12from .subclass_tracking_model import SubclassTrackingModel 

13from .tracking_group import TrackingGroup 

14 

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

16 

17 

18def union( 

19 entity: TrackingGroup | type[SubclassTrackingModel], *, plain: bool | None = None 

20) -> ty.Any: # noqa: ANN401 (type is determined at runtime) 

21 """Get the union of all tracked models in this entity 

22 

23 Parameters 

24 ---------- 

25 entity 

26 The entity for which the union shall be computed. 

27 plain 

28 If set to `True`, a plain union of all members will be returned. 

29 Otherwise, the returned union will be annotated in accordance with 

30 the union mode. 

31 

32 Returns 

33 ------- 

34 Any 

35 If only 1 model is tracked, the model itself will be returned. If >1, 

36 model is tracked a union of all models will be returned. This union may 

37 be an `Annotated` union, depending on the `union_mode` of the entity and 

38 the value of `plain`. 

39 

40 See Also 

41 -------- 

42 [`Union[T]`][dynapydantic.Union] : Wrapper around this function for use 

43 in type annotations. 

44 

45 Raises 

46 ------ 

47 NoRegisteredTypesError 

48 If no models are tracked by this entity. 

49 """ 

50 if isinstance(entity, TrackingGroup): 

51 return entity.union(plain=plain) 

52 if isinstance(entity, type) and issubclass(entity, SubclassTrackingModel): 

53 return entity.__DYNAPYDANTIC__.union(plain=plain) 

54 

55 msg = ( 

56 "dynapydantic.union() works on TrackingGroup or " 

57 f"SubclassTrackingModel, was given {entity}, which was neither." 

58 ) 

59 raise TypeError(msg) 

60 

61 

62def load_plugins(entity: TrackingGroup | type[SubclassTrackingModel]) -> None: 

63 """Load plugins to discover/register additional models 

64 

65 Parameters 

66 ---------- 

67 entity 

68 The entity for which to load plugins 

69 """ 

70 if isinstance(entity, TrackingGroup): 

71 entity.load_plugins() 

72 elif isinstance(entity, type) and issubclass(entity, SubclassTrackingModel): 

73 entity.__DYNAPYDANTIC__.load_plugins() 

74 else: 

75 msg = ( 

76 "dynapydantic.load_plugins() works on TrackingGroup or " 

77 f"SubclassTrackingModel, was given {entity}, which was neither." 

78 ) 

79 raise TypeError(msg) 

80 

81 

82@ty.overload 

83def registered_models(entity: TrackingGroup) -> dict[str, type[pydantic.BaseModel]]: ... 

84 

85 

86@ty.overload 

87def registered_models(entity: type[ModelT]) -> dict[str, type[ModelT]]: ... 

88 

89 

90def registered_models( 

91 entity: TrackingGroup | type[ModelT], 

92) -> dict[str, type[pydantic.BaseModel]] | dict[str, type[ModelT]]: 

93 """Get the mapping of identifier -> model for all models tracked by the entity 

94 

95 Parameters 

96 ---------- 

97 entity 

98 The entity for which to the registered models are desired 

99 

100 Returns 

101 ------- 

102 dict[str, pydantic.BaseModel] 

103 A mapping of identifier to the registered model. This identifier will be 

104 the discriminator value for entities that produce discriminated unions. 

105 If the entity produces a non-discriminated union, the identifier will 

106 just be some unique string. 

107 """ 

108 if isinstance(entity, TrackingGroup): 

109 return entity.models 

110 if isinstance(entity, type) and issubclass(entity, SubclassTrackingModel): 

111 return entity.__DYNAPYDANTIC__.models 

112 

113 msg = ( 

114 "dynapydantic.registered_models() works on TrackingGroup or " 

115 f"SubclassTrackingModel, was given {entity}, which was neither." 

116 ) 

117 raise TypeError(msg)