Coverage for src/kwai/core/json_api.py: 79%
135 statements
« prev ^ index » next coverage.py v7.7.1, created at 2024-01-01 00:00 +0000
1"""Module that defines some JSON:API related models."""
3from typing import Any, Self, cast
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
17class ResourceIdentifier(BaseModel):
18 """A JSON:API resource identifier."""
20 id: str | SkipJsonSchema[None] = None
21 type: str
23 def __hash__(self) -> int:
24 """Create a hash for a resource."""
25 return hash(str(self.id) + self.type)
28class Relationship[R](BaseModel):
29 """A JSON:API relationship."""
31 data: R | list[R] | None
34class ResourceMeta(BaseModel):
35 """Meta for a JSON:API resource."""
37 model_config = ConfigDict(extra="allow")
39 created_at: str
40 updated_at: str | None = None
43class ResourceData[A, R](BaseModel):
44 """A JSON:API resource."""
46 meta: ResourceMeta | SkipJsonSchema[None] = None
47 attributes: A
48 relationships: R | SkipJsonSchema[None] = None
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
61class Meta(BaseModel):
62 """Defines the metadata for the document model.
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)
69 A limit of 0, means there was no limit was set.
70 """
72 model_config = ConfigDict(extra="allow")
74 count: int | None = None
75 offset: int | None = None
76 limit: int | None = None
79class ErrorSource(BaseModel):
80 """Defines the model for an error source."""
82 pointer: str
85class Error(BaseModel):
86 """Defines the model for a JSON:API error."""
88 status: str = ""
89 source: ErrorSource | None = None
90 title: str = ""
91 detail: str = ""
94class BaseDocument[I](BaseModel):
95 """A base model for a JSON:API document."""
97 meta: Meta | SkipJsonSchema[None] = None
98 included: set[I] | SkipJsonSchema[None] = None
99 errors: list[Error] | SkipJsonSchema[None] = None
101 @classmethod
102 def __get_pydantic_json_schema__(
103 cls,
104 core_schema: CoreSchema,
105 handler: GetJsonSchemaHandler,
106 ) -> JsonSchemaValue:
107 """Remove included when T_INCLUDE is NoneType."""
108 json_schema = handler(core_schema)
109 json_schema = handler.resolve_ref_schema(json_schema)
110 if "properties" in json_schema:
111 if "type" in json_schema["properties"]["included"]["items"]:
112 if json_schema["properties"]["included"]["items"]["type"] == "null":
113 del json_schema["properties"]["included"]
114 return json_schema
116 @model_serializer(mode="wrap")
117 def serialize(self, handler) -> dict[str, Any]:
118 """Remove included and meta when the value is None."""
119 result = handler(self)
120 if self.included is None:
121 del result["included"]
122 if self.meta is None:
123 del result["meta"]
124 if not self.errors:
125 del result["errors"]
126 return result
129class SingleDocument[R, I](BaseDocument[I]):
130 """A document that contains only one JSON:API resource."""
132 data: R
134 def __repr__(self):
135 """Return representation of a document."""
136 return f"<{self.__class__.__name__} type={self.data.type}>"
139class MultipleDocument[R, I](BaseDocument[I]):
140 """A document that contains a list of JSON:API resources."""
142 data: list[R] = Field(default_factory=list)
144 def __repr__(self):
145 """Return representation of the document."""
146 if len(self.data) > 0:
147 return f"<{self.__class__.__name__} type={self.data[0].type}[]>"
148 else:
149 return f"<{self.__class__.__name__} type=[]>"
151 def merge(self, other: SingleDocument | Self):
152 """Merge a document into this document.
154 When data is not a list yet, it will be converted to a list. When there are
155 included resources, they will be merged into this document.
156 meta is not merged.
157 """
158 if isinstance(other, SingleDocument):
159 self.data.append(other.data)
160 else:
161 self.data += other.data
162 if other.included is not None:
163 if self.included is None:
164 self.included = other.included
165 else:
166 self.included = self.included.union(other.included)
169class Document[R, I](BaseModel):
170 """A JSON:API document."""
172 meta: Meta | SkipJsonSchema[None] = None
173 data: R | list[R]
174 included: set[I] | SkipJsonSchema[None] = None
175 errors: list[Error] | SkipJsonSchema[None] = None
177 @property
178 def resource(self) -> R:
179 """Return the resource of this document.
181 An assert will occur, when the resource is a list.
182 """
183 assert not isinstance(self.data, list)
184 return self.data
186 @property
187 def resources(self) -> list[R]:
188 """Return the list of resources of this document.
190 An assert will occur, when the resource is not a list.
191 """
192 assert isinstance(self.data, list)
193 return self.data
195 def __repr__(self):
196 """Return representation of a document."""
197 if isinstance(self.data, list):
198 if len(self.data) > 0:
199 return f"<{self.__class__.__name__} type={self.data[0].type}[]>"
200 else:
201 return f"<{self.__class__.__name__} type=[]>"
202 else:
203 return f"<{self.__class__.__name__} type={self.data.type}>"
205 @classmethod
206 def __get_pydantic_json_schema__(
207 cls,
208 core_schema: CoreSchema,
209 handler: GetJsonSchemaHandler,
210 ) -> JsonSchemaValue:
211 """Remove included when T_INCLUDE is NoneType."""
212 json_schema = handler(core_schema)
213 json_schema = handler.resolve_ref_schema(json_schema)
214 if "properties" in json_schema:
215 if "type" in json_schema["properties"]["included"]["items"]:
216 if json_schema["properties"]["included"]["items"]["type"] == "null":
217 del json_schema["properties"]["included"]
218 return json_schema
220 @model_serializer(mode="wrap")
221 def serialize(self, handler) -> dict[str, Any]:
222 """Remove included and meta when the value is None."""
223 result = handler(self)
224 if self.included is None:
225 del result["included"]
226 if self.meta is None:
227 del result["meta"]
228 if not self.errors:
229 del result["errors"]
230 return result
232 def merge(self, other: "Document"):
233 """Merge a document into this document.
235 When data is not a list yet, it will be converted to a list. When there are
236 included resources, they will be merged into this document.
237 meta is not merged.
238 """
239 if not isinstance(self.data, list):
240 self.data = [self.data]
241 self.data.append(cast(Any, other.data))
242 if other.included is not None:
243 if self.included is None:
244 self.included = other.included
245 else:
246 self.included = self.included.union(other.included)
249class PaginationModel(BaseModel):
250 """A model for pagination query parameters.
252 Use this as a dependency on a route. This will handle the page[offset]
253 and page[limit] query parameters.
255 Attributes:
256 offset: The value of the page[offset] query parameter. Default is None.
257 limit: The value of the page[limit] query parameter. Default is None.
258 """
260 offset: int | None = Field(Query(default=None, alias="page[offset]"))
261 limit: int | None = Field(Query(default=None, alias="page[limit]"))
264class JsonApiPresenter[D]:
265 """An interface for a presenter that generates a JSON:API document."""
267 def __init__(self) -> None:
268 self._document: D | None = None
270 def get_document(self) -> D | None:
271 """Return the JSON:API document."""
272 return self._document