Skip to content

Commit 202e0e8

Browse files
committed
chore: upgrade websockets dependency to support 15+ and update import statements
1 parent 7768acb commit 202e0e8

File tree

5 files changed

+207
-94
lines changed

5 files changed

+207
-94
lines changed

.github/workflows/check.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,4 +36,4 @@ jobs:
3636
- name: Generate rest sync code and tests
3737
run: poetry run unasync
3838
- name: Test with pytest
39-
run: poetry run pytest
39+
run: poetry run pytest --verbose --tb=short

ably/transport/websockettransport.py

Lines changed: 28 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,13 @@
1111
from ably.util.eventemitter import EventEmitter
1212
from ably.util.exceptions import AblyException
1313
from ably.util.helper import Timer, unix_time_ms
14-
from websockets.client import WebSocketClientProtocol, connect as ws_connect
14+
try:
15+
# websockets 15+ preferred imports
16+
from websockets import ClientConnection as WebSocketClientProtocol, connect as ws_connect
17+
except ImportError:
18+
# websockets 14 and earlier fallback
19+
from websockets.client import WebSocketClientProtocol, connect as ws_connect
20+
1521
from websockets.exceptions import ConnectionClosedOK, WebSocketException
1622

1723
if TYPE_CHECKING:
@@ -73,24 +79,33 @@ def on_ws_connect_done(self, task: asyncio.Task):
7379

7480
async def ws_connect(self, ws_url, headers):
7581
try:
76-
async with ws_connect(ws_url, extra_headers=headers) as websocket:
77-
log.info(f'ws_connect(): connection established to {ws_url}')
78-
self._emit('connected')
79-
self.websocket = websocket
80-
self.read_loop = self.connection_manager.options.loop.create_task(self.ws_read_loop())
81-
self.read_loop.add_done_callback(self.on_read_loop_done)
82-
try:
83-
await self.read_loop
84-
except WebSocketException as err:
85-
if not self.is_disposed:
86-
await self.dispose()
87-
self.connection_manager.deactivate_transport(err)
82+
# Use additional_headers for websockets 15+, fallback to extra_headers for older versions
83+
try:
84+
async with ws_connect(ws_url, additional_headers=headers) as websocket:
85+
await self._handle_websocket_connection(ws_url, websocket)
86+
except TypeError:
87+
# Fallback for websockets 14 and earlier
88+
async with ws_connect(ws_url, extra_headers=headers) as websocket:
89+
await self._handle_websocket_connection(ws_url, websocket)
8890
except (WebSocketException, socket.gaierror) as e:
8991
exception = AblyException(f'Error opening websocket connection: {e}', 400, 40000)
9092
log.exception(f'WebSocketTransport.ws_connect(): Error opening websocket connection: {exception}')
9193
self._emit('failed', exception)
9294
raise exception
9395

96+
async def _handle_websocket_connection(self, ws_url, websocket):
97+
log.info(f'ws_connect(): connection established to {ws_url}')
98+
self._emit('connected')
99+
self.websocket = websocket
100+
self.read_loop = self.connection_manager.options.loop.create_task(self.ws_read_loop())
101+
self.read_loop.add_done_callback(self.on_read_loop_done)
102+
try:
103+
await self.read_loop
104+
except WebSocketException as err:
105+
if not self.is_disposed:
106+
await self.dispose()
107+
self.connection_manager.deactivate_transport(err)
108+
94109
async def on_protocol_message(self, msg):
95110
self.on_activity()
96111
log.debug(f'WebSocketTransport.on_protocol_message(): received protocol message: {msg}')

0 commit comments

Comments
 (0)