27
27
)
28
28
29
29
try :
30
- from websockets . client import ( # type: ignore[import-not-found,unused-ignore]
31
- WebSocketClientProtocol ,
30
+ from websockets import ( # type: ignore[import-not-found,unused-ignore]
31
+ ClientConnection ,
32
32
connect as ws_connect ,
33
33
)
34
34
from websockets .typing import ( # type: ignore[import-not-found,unused-ignore]
@@ -44,7 +44,7 @@ async def ws_connect(*args, **kwargs): # pylint: disable=unused-argument
44
44
raise NotImplementedError ("Subscriptions require 'websockets' package." )
45
45
yield # pylint: disable=unreachable
46
46
47
- WebSocketClientProtocol = Any # type: ignore[misc,assignment,unused-ignore]
47
+ ClientConnection = Any # type: ignore[misc,assignment,unused-ignore]
48
48
Data = Any # type: ignore[misc,assignment,unused-ignore]
49
49
Origin = Any # type: ignore[misc,assignment,unused-ignore]
50
50
@@ -392,7 +392,7 @@ async def _execute_ws(
392
392
if data :
393
393
yield data
394
394
395
- async def _send_connection_init (self , websocket : WebSocketClientProtocol ) -> None :
395
+ async def _send_connection_init (self , websocket : ClientConnection ) -> None :
396
396
payload : Dict [str , Any ] = {
397
397
"type" : GraphQLTransportWSMessageType .CONNECTION_INIT .value
398
398
}
@@ -402,7 +402,7 @@ async def _send_connection_init(self, websocket: WebSocketClientProtocol) -> Non
402
402
403
403
async def _send_subscribe (
404
404
self ,
405
- websocket : WebSocketClientProtocol ,
405
+ websocket : ClientConnection ,
406
406
operation_id : str ,
407
407
query : str ,
408
408
operation_name : Optional [str ] = None ,
@@ -422,7 +422,7 @@ async def _send_subscribe(
422
422
async def _handle_ws_message (
423
423
self ,
424
424
message : Data ,
425
- websocket : WebSocketClientProtocol ,
425
+ websocket : ClientConnection ,
426
426
expected_type : Optional [GraphQLTransportWSMessageType ] = None ,
427
427
) -> Optional [Dict [str , Any ]]:
428
428
try :
@@ -601,7 +601,7 @@ async def _execute_ws_with_telemetry(
601
601
yield data
602
602
603
603
async def _send_connection_init_with_telemetry (
604
- self , root_span : Span , websocket : WebSocketClientProtocol
604
+ self , root_span : Span , websocket : ClientConnection
605
605
) -> None :
606
606
with self .tracer .start_as_current_span ( # type: ignore
607
607
"connection init" , context = set_span_in_context (root_span )
@@ -620,7 +620,7 @@ async def _send_connection_init_with_telemetry(
620
620
async def _send_subscribe_with_telemetry (
621
621
self ,
622
622
root_span : Span ,
623
- websocket : WebSocketClientProtocol ,
623
+ websocket : ClientConnection ,
624
624
operation_id : str ,
625
625
query : str ,
626
626
operation_name : Optional [str ] = None ,
@@ -652,7 +652,7 @@ async def _handle_ws_message_with_telemetry(
652
652
self ,
653
653
root_span : Span ,
654
654
message : Data ,
655
- websocket : WebSocketClientProtocol ,
655
+ websocket : ClientConnection ,
656
656
expected_type : Optional [GraphQLTransportWSMessageType ] = None ,
657
657
) -> Optional [Dict [str , Any ]]:
658
658
with self .tracer .start_as_current_span ( # type: ignore
0 commit comments