1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
|
from aioitertools.itertools import tee as atee
from starlette.requests import Request
from starlette.responses import Response
from openapi_core.contrib.starlette.requests import StarletteOpenAPIRequest
from openapi_core.contrib.starlette.responses import StarletteOpenAPIResponse
from openapi_core.unmarshalling.processors import AsyncUnmarshallingProcessor
from openapi_core.unmarshalling.typing import ErrorsHandlerCallable
class StarletteIntegration(AsyncUnmarshallingProcessor[Request, Response]):
request_cls = StarletteOpenAPIRequest
response_cls = StarletteOpenAPIResponse
async def get_openapi_request(
self, request: Request
) -> StarletteOpenAPIRequest:
body = await request.body()
return self.request_cls(request, body)
async def get_openapi_response(
self, response: Response
) -> StarletteOpenAPIResponse:
assert self.response_cls is not None
data = None
if hasattr(response, "body_iterator"):
body_iter1, body_iter2 = atee(response.body_iterator)
response.body_iterator = body_iter2
data = b"".join(
[
(
chunk.encode(response.charset)
if not isinstance(chunk, bytes)
else chunk
)
async for chunk in body_iter1
]
)
return self.response_cls(response, data=data)
def should_validate_response(self) -> bool:
return self.response_cls is not None
async def handle_response(
self,
request: Request,
response: Response,
errors_handler: ErrorsHandlerCallable[Response],
) -> Response:
if not self.should_validate_response():
return response
return await super().handle_response(request, response, errors_handler)
|
