Skip to content

Commit

Permalink
Simplify code.
Browse files Browse the repository at this point in the history
It had to be written in that way with asyncio.wait_for but that isn't
necessary anymore with asyncio.timeout.
  • Loading branch information
aaugustin committed Jan 13, 2024
1 parent 3c6b1aa commit dbb745b
Showing 1 changed file with 24 additions and 27 deletions.
51 changes: 24 additions & 27 deletions src/websockets/legacy/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -644,34 +644,31 @@ def __await__(self) -> Generator[Any, None, WebSocketClientProtocol]:

async def __await_impl_timeout__(self) -> WebSocketClientProtocol:
async with asyncio_timeout(self.open_timeout):
return await self.__await_impl__()

async def __await_impl__(self) -> WebSocketClientProtocol:
for redirects in range(self.MAX_REDIRECTS_ALLOWED):
_transport, _protocol = await self._create_connection()
protocol = cast(WebSocketClientProtocol, _protocol)
try:
await protocol.handshake(
self._wsuri,
origin=protocol.origin,
available_extensions=protocol.available_extensions,
available_subprotocols=protocol.available_subprotocols,
extra_headers=protocol.extra_headers,
)
except RedirectHandshake as exc:
protocol.fail_connection()
await protocol.wait_closed()
self.handle_redirect(exc.uri)
# Avoid leaking a connected socket when the handshake fails.
except (Exception, asyncio.CancelledError):
protocol.fail_connection()
await protocol.wait_closed()
raise
for _redirects in range(self.MAX_REDIRECTS_ALLOWED):
_transport, _protocol = await self._create_connection()
protocol = cast(WebSocketClientProtocol, _protocol)
try:
await protocol.handshake(
self._wsuri,
origin=protocol.origin,
available_extensions=protocol.available_extensions,
available_subprotocols=protocol.available_subprotocols,
extra_headers=protocol.extra_headers,
)
except RedirectHandshake as exc:
protocol.fail_connection()
await protocol.wait_closed()
self.handle_redirect(exc.uri)
# Avoid leaking a connected socket when the handshake fails.
except (Exception, asyncio.CancelledError):
protocol.fail_connection()
await protocol.wait_closed()
raise
else:
self.protocol = protocol
return protocol
else:
self.protocol = protocol
return protocol
else:
raise SecurityError("too many redirects")
raise SecurityError("too many redirects")

# ... = yield from connect(...) - remove when dropping Python < 3.10

Expand Down

0 comments on commit dbb745b

Please sign in to comment.