Coverage for src/kwai/api/v1/trainings/schemas/training.py: 92%

1"""Schemas for training(s).""" 

2 

3from types import NoneType 

4from typing import Annotated 

5 

6from pydantic import BaseModel, Field 

7 

8from kwai.api.converter import MarkdownConverter 

9from kwai.api.schemas.resources import ( 

10 CoachResourceIdentifier, 

11 TeamResourceIdentifier, 

12 TrainingDefinitionResourceIdentifier, 

13 TrainingResourceIdentifier, 

14) 

15from kwai.api.v1.trainings.schemas.team import TeamDocument, TeamResource 

16from kwai.api.v1.trainings.schemas.training_definition import ( 

17 TrainingDefinitionDocument, 

18 TrainingDefinitionResource, 

19) 

20from kwai.core.json_api import Document, Relationship, ResourceData 

21from kwai.modules.training.trainings.training import TrainingEntity 

22 

23 

24class TrainingText(BaseModel): 

25 """Schema for the content of a training.""" 

26 

27 locale: str 

28 format: str 

29 title: str 

30 summary: str 

31 content: str | None 

32 original_summary: str | None 

33 original_content: str | None 

34 

35 

36class TrainingCoach(BaseModel): 

37 """Schema for coach/training specific information.""" 

38 

39 id: str 

40 head: bool 

41 present: bool 

42 payed: bool 

43 

44 

45class TrainingEvent(BaseModel): 

46 """Schema for the event information of a training.""" 

47 

48 start_date: str 

49 end_date: str 

50 location: str 

51 cancelled: bool 

52 active: bool 

53 

54 

55class CoachAttributes(BaseModel): 

56 """Attributes for a coach JSON:API resource.""" 

57 

58 name: str 

59 

60 

61class CoachResource(CoachResourceIdentifier, ResourceData[CoachAttributes, NoneType]): 

62 """A JSON:API resource for a coach.""" 

63 

64 

65class TrainingAttributes(BaseModel): 

66 """Attributes for training JSON:API resource.""" 

67 

68 texts: list[TrainingText] 

69 event: TrainingEvent 

70 remark: str 

71 coaches: list[TrainingCoach] 

72 

73 

74class TrainingRelationships(BaseModel): 

75 """Relationships of a training JSON:API resource.""" 

76 

77 coaches: Relationship[CoachResourceIdentifier] 

78 teams: Relationship[TeamResourceIdentifier] 

79 definition: Relationship[TrainingDefinitionResourceIdentifier] 

80 

81 

82class TrainingResource( 

83 TrainingResourceIdentifier, ResourceData[TrainingAttributes, TrainingRelationships] 

84): 

85 """A JSON:API resource for a training.""" 

86 

87 

88TrainingInclude = Annotated[ 

89 TeamResource | TrainingDefinitionResource | CoachResource, 

90 Field(discriminator="type"), 

91] 

92 

93 

94class TrainingDocument(Document[TrainingResource, TrainingInclude]): 

95 """A JSON:API document for one or more training resources.""" 

96 

97 @classmethod 

98 def create(cls, training: TrainingEntity) -> "TrainingDocument": 

99 """Create a training document from a training entity.""" 

100 training_resource = TrainingResource( 

101 id=str(training.id), 

102 attributes=TrainingAttributes( 

103 texts=[ 

104 TrainingText( 

105 locale=text.locale.value, 

106 format=text.format.value, 

107 title=text.title, 

108 summary=MarkdownConverter().convert(text.summary), 

109 content=MarkdownConverter().convert(text.content) 

110 if text.content 

111 else None, 

112 original_summary=text.summary, 

113 original_content=text.content, 

114 ) 

115 for text in training.texts 

116 ], 

117 event=TrainingEvent( 

118 start_date=str(training.period.start_date), 

119 end_date=str(training.period.end_date), 

120 location=training.location or "", 

121 cancelled=training.cancelled, 

122 active=training.active, 

123 ), 

124 remark=training.remark or "", 

125 coaches=[ 

126 TrainingCoach( 

127 id=str(coach.coach.id), 

128 head=coach.type == 1, 

129 present=coach.present, 

130 payed=coach.payed, 

131 ) 

132 for coach in training.coaches 

133 ], 

134 ), 

135 ) 

136 

137 included: set[TrainingInclude] = set() 

138 

139 training_resource.relationships = TrainingRelationships( 

140 coaches=Relationship[CoachResourceIdentifier](data=[]), 

141 teams=Relationship[TeamResourceIdentifier](data=[]), 

142 definition=Relationship[TrainingDefinitionResourceIdentifier](data=None), 

143 ) 

144 

145 for training_coach in training.coaches: 

146 training_resource.relationships.coaches.data.append( 

147 CoachResourceIdentifier(id=str(training_coach.coach.id)) 

148 ) 

149 included.add( 

150 CoachResource( 

151 id=str(training_coach.coach.id), 

152 attributes=CoachAttributes(name=str(training_coach.coach.name)), 

153 ) 

154 ) 

155 if training.definition: 

156 training_resource.relationships.definition = Relationship[ 

157 TrainingDefinitionResourceIdentifier 

158 ]( 

159 data=( 

160 TrainingDefinitionResourceIdentifier(id=str(training.definition.id)) 

161 ) 

162 ) 

163 training_definition_document = TrainingDefinitionDocument.create( 

164 training.definition 

165 ) 

166 included.add(training_definition_document.data) 

167 included = included.union(training_definition_document.included) 

168 

169 for team in training.teams: 

170 training_resource.relationships.teams.data.append( 

171 TeamResourceIdentifier(id=str(team.id)) 

172 ) 

173 team_document = TeamDocument.create(team) 

174 included.add(team_document.data) 

175 

176 return TrainingDocument(data=training_resource, included=included)