class ASGIContext:
_ws_close_codes = frozenset((
WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED, WSMsgType.ERROR,
))
def __init__(
self, app: t.Callable[..., t.Any],
request: Request, root_path: str,
):
self.request = request
connection_hdr = request.headers.get("Connection", "").lower()
self.http_version = "1.1" if connection_hdr == 'keep-alive' else '1.0'
self.app = app
self.root_path = root_path.rstrip("/")
self.start_response_event = asyncio.Event()
self.ws_connect_event = asyncio.Event()
self.response = None # type: _ResponseType
self.writer = None # type: _WriterType
self.task = None # type: t.Optional[asyncio.Task]
self.loop = asyncio.get_event_loop()
def is_websocket(self):
return (
self.request.headers.get("Connection", "").lower() == "upgrade" and
self.request.headers.get("Upgrade", "").lower() == "websocket"
)
@property
def scope(self) -> dict:
raw_path = self.request.raw_path
result = {
"type": "http",
"http_version": self.http_version,
"server": [self.request.url.host, self.request.url.port],
"client": [self.request.remote, 0],
"scheme": self.request.url.scheme,
"method": self.request.method,
"root_path": self.root_path,
"path": self.request.path,
"raw_path": raw_path.encode(),
"query_string": self.request.query_string.encode(),
"headers": [hdr for hdr in self.request.raw_headers],
}
if self.is_websocket():
result["type"] = "websocket"
result["scheme"] = "wss" if self.request.secure else "ws"
result["subprotocols"] = []
return result
async def on_receive(self):
if self.is_websocket():
if not self.ws_connect_event.is_set():
self.ws_connect_event.set()
return {
"type": "websocket.connect",
"headers": tuple(self.request.raw_headers),
}
while True:
msg = await self.response.receive()
if msg.type in (WSMsgType.BINARY, WSMsgType.TEXT):
bytes_payload = None
str_payload = None
if msg.type == WSMsgType.BINARY:
bytes_payload = msg.data
if msg.type == WSMsgType.TEXT:
str_payload = msg.data
return {
"type": "websocket.receive",
"bytes": bytes_payload,
"text": str_payload,
}
if msg.type in self._ws_close_codes:
self.start_response_event.set()
return {
"type": "websocket.disconnect",
"code": self.response.close_code,
}
chunk, more_body = await self.request.content.readchunk()
return {
"type": "http.request",
"body": chunk,
"more_body": more_body,
}
async def on_send(self, payload: t.Dict[str, t.Any]):
if payload["type"] == "http.response.start":
if self.start_response_event.is_set():
raise asyncio.InvalidStateError
self.response = StreamResponse()
self.response.set_status(payload["status"])
for name, value in payload.get("headers", ()):
header_name = name.title().decode()
self.response.headers[header_name] = value.decode()
self.writer = await self.response.prepare(self.request)
self.start_response_event.set()
return
if payload["type"] == "websocket.accept":
if self.start_response_event.is_set():
raise asyncio.InvalidStateError
self.response = WebSocketResponse()
self.writer = await self.response.prepare(self.request)
return
if payload["type"] == "http.response.body":
if self.writer is None:
raise TypeError("Unexpected message %r" % payload, payload)
await self.writer.write(payload["body"])
# receive_queue.put_nowait({"type": "http.disconnect"})
if payload.get("more_body", False):
await self.writer.write_eof()
return
if payload["type"] == "websocket.send":
if (
isinstance(self.response, WebSocketResponse) and
self.response.closed
):
raise TypeError("Unexpected message %r" % payload, payload)
if not isinstance(self.response, WebSocketResponse):
raise RuntimeError("Wrong response type")
message_bytes = payload.get("bytes")
message_text = payload.get("text")
if not any((message_text, message_bytes)):
raise TypeError(
"Exactly one of bytes or text must be non-None."
" One or both keys may be present, however.",
)
if message_bytes is not None:
await self.response.send_bytes(message_bytes)
if message_text is not None:
await self.response.send_str(message_text)
return
async def get_response(self) -> t.Union[StreamResponse, WebSocketResponse]:
self.task = self.loop.create_task(
self.app(
self.scope,
self.on_receive,
self.on_send,
),
)
try:
await self.start_response_event.wait()
except asyncio.CancelledError:
if not self.task.done():
self.task.cancel()
raise
if self.response is None:
raise RuntimeError
return self.response
