Coverage for src/kwai/core/json_api.py: 86%

1"""Module that defines some JSON:API related models.""" 

2 

3from typing import Any, cast 

4 

5from fastapi import Query 

6from pydantic import ( 

7 BaseModel, 

8 ConfigDict, 

9 Field, 

10 GetJsonSchemaHandler, 

11 model_serializer, 

12) 

13from pydantic.json_schema import JsonSchemaValue, SkipJsonSchema 

14from pydantic_core import CoreSchema 

15 

16 

17class ResourceIdentifier(BaseModel): 

18 """A JSON:API resource identifier.""" 

19 

20 id: str | SkipJsonSchema[None] = None 

21 type: str 

22 

23 def __hash__(self) -> int: 

24 """Create a hash for a resource.""" 

25 return hash(str(self.id) + self.type) 

26 

27 

28class Relationship[R](BaseModel): 

29 """A JSON:API relationship.""" 

30 

31 data: R | list[R] | None 

32 

33 

34class ResourceMeta(BaseModel): 

35 """Meta for a JSON:API resource.""" 

36 

37 model_config = ConfigDict(extra="allow") 

38 

39 created_at: str 

40 updated_at: str | None = None 

41 

42 

43class ResourceData[A, R](BaseModel): 

44 """A JSON:API resource.""" 

45 

46 meta: ResourceMeta | SkipJsonSchema[None] = None 

47 attributes: A 

48 relationships: R | SkipJsonSchema[None] = None 

49 

50 @model_serializer(mode="wrap") 

51 def serialize(self, handler) -> dict[str, Any]: 

52 """Remove relationships and meta from serialization when the values are none.""" 

53 result = handler(self) 

54 if self.relationships is None: 

55 del result["relationships"] 

56 if self.meta is None: 

57 del result["meta"] 

58 return result 

59 

60 

61class Meta(BaseModel): 

62 """Defines the metadata for the document model. 

63 

64 Attributes: 

65 count: The number of actual resources 

66 offset: The offset of the returned resources (pagination) 

67 limit: The maximum number of returned resources (pagination) 

68 

69 A limit of 0, means there was no limit was set. 

70 """ 

71 

72 model_config = ConfigDict(extra="allow") 

73 

74 count: int | None = None 

75 offset: int | None = None 

76 limit: int | None = None 

77 

78 

79class ErrorSource(BaseModel): 

80 """Defines the model for an error source.""" 

81 

82 pointer: str 

83 

84 

85class Error(BaseModel): 

86 """Defines the model for a JSON:API error.""" 

87 

88 status: str = "" 

89 source: ErrorSource | None = None 

90 title: str = "" 

91 detail: str = "" 

92 

93 

94class Document[R, I](BaseModel): 

95 """A JSON:API document.""" 

96 

97 meta: Meta | SkipJsonSchema[None] = None 

98 data: R | list[R] 

99 included: set[I] | SkipJsonSchema[None] = None 

100 errors: list[Error] | SkipJsonSchema[None] = None 

101 

102 @property 

103 def resource(self) -> R: 

104 """Return the resource of this document. 

105 

106 An assert will occur, when the resource is a list. 

107 """ 

108 assert not isinstance(self.data, list) 

109 return self.data 

110 

111 @property 

112 def resources(self) -> list[R]: 

113 """Return the list of resources of this document. 

114 

115 An assert will occur, when the resource is not a list. 

116 """ 

117 assert isinstance(self.data, list) 

118 return self.data 

119 

120 def __repr__(self): 

121 """Return representation of a document.""" 

122 if isinstance(self.data, list): 

123 if len(self.data) > 0: 

124 return f"<{self.__class__.__name__} type={self.data[0].type}[]>" 

125 else: 

126 return f"<{self.__class__.__name__} type=[]>" 

127 else: 

128 return f"<{self.__class__.__name__} type={self.data.type}>" 

129 

130 @classmethod 

131 def __get_pydantic_json_schema__( 

132 cls, 

133 core_schema: CoreSchema, 

134 handler: GetJsonSchemaHandler, 

135 ) -> JsonSchemaValue: 

136 """Remove included when T_INCLUDE is NoneType.""" 

137 json_schema = handler(core_schema) 

138 json_schema = handler.resolve_ref_schema(json_schema) 

139 if "properties" in json_schema: 

140 if "type" in json_schema["properties"]["included"]["items"]: 

141 if json_schema["properties"]["included"]["items"]["type"] == "null": 

142 del json_schema["properties"]["included"] 

143 return json_schema 

144 

145 @model_serializer(mode="wrap") 

146 def serialize(self, handler) -> dict[str, Any]: 

147 """Remove included and meta when the value is None.""" 

148 result = handler(self) 

149 if self.included is None: 

150 del result["included"] 

151 if self.meta is None: 

152 del result["meta"] 

153 if not self.errors: 

154 del result["errors"] 

155 return result 

156 

157 def merge(self, other: "Document"): 

158 """Merge a document into this document. 

159 

160 When data is not a list yet, it will be converted to a list. When there are 

161 included resources, they will be merged into this document. 

162 meta is not merged. 

163 """ 

164 if not isinstance(self.data, list): 

165 self.data = [self.data] 

166 self.data.append(cast(Any, other.data)) 

167 if other.included is not None: 

168 if self.included is None: 

169 self.included = other.included 

170 else: 

171 self.included = self.included.union(other.included) 

172 

173 

174class PaginationModel(BaseModel): 

175 """A model for pagination query parameters. 

176 

177 Use this as a dependency on a route. This will handle the page[offset] 

178 and page[limit] query parameters. 

179 

180 Attributes: 

181 offset: The value of the page[offset] query parameter. Default is None. 

182 limit: The value of the page[limit] query parameter. Default is None. 

183 """ 

184 

185 offset: int | None = Field(Query(default=None, alias="page[offset]")) 

186 limit: int | None = Field(Query(default=None, alias="page[limit]")) 

187 

188 

189class JsonApiPresenter[Document]: 

190 """An interface for a presenter that generates a JSON:API document.""" 

191 

192 def __init__(self): 

193 self._document: Document = None 

194 

195 def get_document(self) -> Document: 

196 """Return the JSON:API document.""" 

197 return self._document