class WebsocketHandshake(asyncio.Protocol):
def __init__(self, client, loop):
self.client = client
self.loop = loop
self.transport: asyncio.Transport = None
self.parser = HttpResponseParser(self)
self.current_status = None
self.current_headers = None
def connection_made(self, transport):
"""
:param transport:
:return:
"""
wr = WebsocketRequest(self.client.host,
path=self.client.path,
origin=self.client.origin)
transport.write(wr.encode())
self.transport = transport
print('connected')
def data_received(self, data):
self.parser.feed(data)
print(f'Data received: {data}')
def connection_lost(self, exc):
print('The server closed the connection')
print('Stop the event loop')
# Parser Callbacks
def on_body(self):
pass
def on_headers_complete(self, headers, status_code):
self.current_status = status_code
self.current_headers = headers
def on_message_complete(self):
self.transport.set_protocol(
WebsocketProtocol(self.transport, self.loop))
class Response:
__slots__ = (
"_connection",
"_headers",
"_content",
"_parser",
"_parser_status",
"_cookies",
"_status_code",
"_decode",
"_chunk_size",
"_decoder",
"_encoding",
"request",
"url",
)
def __init__(
self,
url: str,
connection: Connection,
request: Request,
chunk_size: int = 1 * 1024 * 1024,
decode: bool = True,
):
self._connection = connection
self._headers = {}
self._content = bytearray()
# noinspection PyArgumentList
self._parser = HttpResponseParser(self)
self._parser_status = ResponseStatus.PENDING_HEADERS
self._cookies = None
self._status_code = None
self._chunk_size = chunk_size
self._decode = decode
self._decoder = None
self._encoding = None
self.request = request
self.url = url
@property
def encoding(self):
return "utf-8"
def json(self, *args, loads=None, **kwargs):
if not loads:
loads = json.loads
return loads(self.content.decode(self.encoding), *args, **kwargs)
def text(self, encoding: str = None) -> str:
return self.content.decode(encoding=encoding or self.encoding)
@property
def status_code(self):
"""
:return:
"""
if self._parser_status == ResponseStatus.PENDING_HEADERS:
raise Exception(
"Status code not loaded yet. "
"In streaming mode you should manually call load_headers().")
return self._status_code
@property
def headers(self):
"""
:return:
"""
if self._parser_status == ResponseStatus.PENDING_HEADERS:
raise Exception(
"Headers not loaded yet. "
"In streaming mode you should manually call load_headers().")
return self._headers
@property
def content(self):
"""
:return:
"""
if self._parser_status == ResponseStatus.PENDING_BODY:
raise Exception("You need to call read_content() "
"before using this in streaming mode.")
return self._content
@property
async def cookies(self):
"""
:return:
"""
if self._parser_status == ResponseStatus.PENDING_HEADERS:
await self.receive_headers()
if self._cookies is None:
self._cookies = CookiesJar()
if self._headers.get("set-cookie"):
self._cookies.add_cookie(
Cookie.from_header(self._headers["set-cookie"]))
return self._cookies
async def read_content(self):
"""
:return:
"""
await self.receive_headers()
if self._parser_status != ResponseStatus.PENDING_BODY:
return
try:
length = int(self._headers["Content-Length"])
if not self._decoder:
# Skipping the HTTP parser for performance.
self._content.extend(await
self._connection.read_exactly(length))
self.on_message_complete()
else:
self._parser.feed(await self._connection.read_exactly(length))
except KeyError:
if self._headers.get("Transfer-Encoding") == "chunked":
await self._handle_chunked_encoding()
else:
raise Exception("Invalid response.")
return self._content
except ValueError:
raise Exception("Invalid content-length header.")
async def _handle_chunked_encoding(self):
"""
:return:
"""
self._parser_status = ResponseStatus.CHUNKED_TRANSFER
while self._parser_status == ResponseStatus.CHUNKED_TRANSFER:
self._parser.feed(await self._connection.read_until(b"\r\n"))
async def receive_headers(self):
"""
:return:
"""
if self._parser_status != ResponseStatus.PENDING_HEADERS:
return
try:
self._parser.feed(await self._connection.read_until(b"\r\n\r\n"))
if self._headers.get("content-encoding") == "gzip":
self._decoder = GzipDecoder()
self._parser_status = ResponseStatus.PENDING_BODY
except OSError as error:
# In case any network error occurs we should discard this connection.
self._parser_status = ResponseStatus.ERROR
self._connection.close()
raise error
async def _release_connection(self):
"""
:return:
"""
await self._connection.pool.release_connection(
self._connection, self._parser.should_keep_alive())
async def stream(
self,
chunk_size: int = 1 * 1024 * 1024,
chunk_timeout: int = 10,
complete_timeout: int = 300,
):
"""
:param complete_timeout:
:param chunk_timeout:
:param chunk_size:
:return:
"""
if self._parser_status not in (
ResponseStatus.PENDING_HEADERS,
ResponseStatus.PENDING_BODY,
):
raise StreamAlreadyConsumed
if self._parser_status == ResponseStatus.PENDING_HEADERS:
await wait_for(self.receive_headers(), chunk_timeout)
try:
length = int(self.headers["Content-Length"])
remaining = length
while remaining:
bytes_to_read = min(remaining, chunk_size)
task = self._connection.read_exactly(bytes_to_read)
start_time = time.time()
self._parser.feed(await wait_for(
task, min(chunk_timeout, complete_timeout)))
complete_timeout -= time.time() - start_time
remaining -= bytes_to_read
yield bytes(self._content)
self._content = bytearray()
except KeyError:
if self._headers.get("Transfer-Encoding") == "chunked":
self._parser_status = ResponseStatus.CHUNKED_TRANSFER
while self._parser_status == ResponseStatus.CHUNKED_TRANSFER:
task = self._connection.read_until(b"\r\n")
start_time = time.time()
try:
self._parser.feed(await wait_for(
task, min(chunk_timeout, complete_timeout)))
except asyncio.LimitOverrunError as error:
self._parser.feed(await self._connection.read_exactly(
error.consumed))
complete_timeout -= time.time() - start_time
while len(self._content) >= chunk_size:
yield self._content[:chunk_size]
self._content = self._content[chunk_size:]
while len(self._content) > 0:
yield self._content[:chunk_size]
self._content = self._content[chunk_size:]
else:
raise Exception("Invalid response.")
except ValueError:
raise Exception("Invalid content-length header.")
def is_redirect(self):
return self.headers.get("location") is not None
#############
# # # Http Parser Callbacks
######################################
def on_headers_complete(self, headers, status_code):
"""
:param headers:
:param status_code:
:return:
"""
self._headers = headers
self._status_code = status_code
def on_body(self, data: bytes):
"""
:param data:
:return:
"""
if self._decoder is not None:
self._content.extend(self._decoder.decompress(data))
else:
self._content.extend(data)
def on_message_complete(self):
"""
:return:
"""
self._parser_status = ResponseStatus.COMPLETE
try:
self._connection.loop.create_task(self._release_connection())
except Exception as error:
print(error)
def chunk_complete(self):
pass
def __repr__(self):
if 400 > self.status_code >= 300:
if len(self.url) > 30:
url = self.url[:30] + "..."
else:
url = self.url
return f'<Response [{self.status_code} => "{url}"]>'
else:
return f"<Response [{self.status_code}]>"