flo-token-explorer-historic/py3/lib/python3.6/site-packages/quart/asgi.py
2021-12-09 18:16:18 +00:00

190 lines
6.8 KiB
Python

import asyncio
from functools import partial
from typing import AnyStr, Callable, Set, TYPE_CHECKING
from .datastructures import CIMultiDict
from .wrappers import Request, Response, Websocket # noqa: F401
if TYPE_CHECKING:
from .app import Quart # noqa: F401
class ASGIHTTPConnection:
def __init__(self, app: 'Quart', scope: dict) -> None:
self.app = app
self.scope = scope
async def __call__(self, receive: Callable, send: Callable) -> None:
request = self._create_request_from_scope()
receiver_task = asyncio.ensure_future(self.handle_messages(request, receive))
handler_task = asyncio.ensure_future(self.handle_request(request, send))
_, pending = await asyncio.wait(
[handler_task, receiver_task], return_when=asyncio.FIRST_COMPLETED,
)
await _cancel_tasks(pending)
async def handle_messages(self, request: Request, receive: Callable) -> None:
while True:
message = await receive()
if message['type'] == 'http.request':
request.body.append(message['body'])
if not message.get('more_body', False):
request.body.set_complete()
elif message['type'] == 'http.disconnect':
return
def _create_request_from_scope(self) -> Request:
headers = CIMultiDict()
headers['Remote-Addr'] = (self.scope.get('client') or ['<local>'])[0]
for name, value in self.scope['headers']:
headers.add(name.decode().title(), value.decode())
if self.scope['http_version'] < '1.1':
headers.setdefault('Host', self.app.config['SERVER_NAME'] or '')
return self.app.request_class(
self.scope['method'], self.scope['scheme'], self.scope['path'],
self.scope['query_string'], headers,
max_content_length=self.app.config['MAX_CONTENT_LENGTH'],
body_timeout=self.app.config['BODY_TIMEOUT'],
)
async def handle_request(self, request: Request, send: Callable) -> None:
response = await self.app.handle_request(request)
try:
await asyncio.wait_for(self._send_response(send, response), timeout=response.timeout)
except asyncio.TimeoutError:
pass
async def _send_response(self, send: Callable, response: Response) -> None:
headers = [
(key.lower().encode(), value.encode())
for key, value in response.headers.items()
]
await send({
'type': 'http.response.start',
'status': response.status_code,
'headers': headers,
})
if 'http.response.push' in self.scope.get('extensions', {}):
for path in response.push_promises:
await send({
'type': 'http.response.push',
'path': path,
'headers': [],
})
async for data in response.response:
await send({
'type': 'http.response.body',
'body': data,
'more_body': True,
})
await send({
'type': 'http.response.body',
'body': b'',
'more_body': False,
})
class ASGIWebsocketConnection:
def __init__(self, app: 'Quart', scope: dict) -> None:
self.app = app
self.scope = scope
self.queue: asyncio.Queue = asyncio.Queue()
self._accepted = False
async def __call__(self, receive: Callable, send: Callable) -> None:
websocket = self._create_websocket_from_scope(send)
receiver_task = asyncio.ensure_future(self.handle_messages(receive))
handler_task = asyncio.ensure_future(self.handle_websocket(websocket, send))
_, pending = await asyncio.wait(
[handler_task, receiver_task], return_when=asyncio.FIRST_COMPLETED,
)
await _cancel_tasks(pending)
async def handle_messages(self, receive: Callable) -> None:
while True:
event = await receive()
if event['type'] == 'websocket.receive':
await self.queue.put(event.get('bytes') or event['text'])
elif event['type'] == 'websocket.disconnect':
return
def _create_websocket_from_scope(self, send: Callable) -> Websocket:
headers = CIMultiDict()
headers['Remote-Addr'] = (self.scope.get('client') or ['<local>'])[0]
for name, value in self.scope['headers']:
headers.add(name.decode().title(), value.decode())
return self.app.websocket_class(
self.scope['path'], self.scope['query_string'], self.scope['scheme'],
headers, self.queue, partial(self.send_data, send),
partial(self.accept_connection, send),
)
async def handle_websocket(self, websocket: Websocket, send: Callable) -> None:
response = await self.app.handle_websocket(websocket)
if (
response is not None and not self._accepted
and 'websocket.http.response' in self.scope.get('extensions', {})
):
headers = [
(key.lower().encode(), value.encode())
for key, value in response.headers.items()
]
await send({
'type': 'websocket.http.response.start',
'status': response.status_code,
'headers': headers,
})
async for data in response.response:
await send({
'type': 'websocket.http.response.body',
'body': data,
'more_body': True,
})
await send({
'type': 'websocket.http.response.body',
'body': b'',
'more_body': False,
})
elif self._accepted:
await send({
'type': 'websocket.close',
'code': 1000,
})
async def send_data(self, send: Callable, data: AnyStr) -> None:
if isinstance(data, str):
await send({
'type': 'websocket.send',
'text': data,
})
else:
await send({
'type': 'websocket.send',
'bytes': data,
})
async def accept_connection(self, send: Callable) -> None:
if not self._accepted:
await send({
'type': 'websocket.accept',
})
self._accepted = True
async def _cancel_tasks(tasks: Set[asyncio.Future]) -> None:
# Cancel any pending, and wait for the cancellation to
# complete i.e. finish any remaining work.
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
# Raise any unexcepted exceptions
for task in tasks:
if not task.cancelled() and task.exception() is not None:
raise task.exception()