From 7b522ec0df8f4e26abe09046a0ae7861714f5a2a Mon Sep 17 00:00:00 2001
From: Aymeric Augustin <aymeric.augustin@fractalideas.com>
Date: Sat, 13 Jan 2024 21:54:08 +0100
Subject: [PATCH] Simplify code.

It had to be written in that way with asyncio.wait_for but that isn't
necessary anymore with asyncio.timeout.
---
 src/websockets/legacy/client.py | 55 ++++++++++++++++-----------------
 1 file changed, 26 insertions(+), 29 deletions(-)

diff --git a/src/websockets/legacy/client.py b/src/websockets/legacy/client.py
index 48622523..b85d2286 100644
--- a/src/websockets/legacy/client.py
+++ b/src/websockets/legacy/client.py
@@ -640,38 +640,35 @@ async def __aexit__(
 
     def __await__(self) -> Generator[Any, None, WebSocketClientProtocol]:
         # Create a suitable iterator by calling __await__ on a coroutine.
-        return self.__await_impl_timeout__().__await__()
-
-    async def __await_impl_timeout__(self) -> WebSocketClientProtocol:
-        async with asyncio_timeout(self.open_timeout):
-            return await self.__await_impl__()
+        return self.__await_impl__().__await__()
 
     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
+        async with asyncio_timeout(self.open_timeout):
+            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